Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- # set up Python environment: numpy for numerical routines, and matplotlib for plotting
- import time
- import numpy as np
- import matplotlib.pyplot as plt
- # display plots in this notebook
- # set display defaults
- plt.rcParams['figure.figsize'] = (10, 10) # large images
- plt.rcParams['image.interpolation'] = 'nearest' # don't interpolate: show square pixels
- plt.rcParams['image.cmap'] = 'gray' # use grayscale output rather than a (potentially misleading) color heatmap
- import os
- import sys
- import caffe
- caffe.set_mode_cpu()
- start_time = time.time()
- net_root = 'D:\\Development\\caffe-windows\\fcn.berkeleyvision.org-master\\voc-fcn8s'
- model_def = net_root + '\\deploy.prototxt'
- model_weights = net_root + '\\fcn8s-heavy-40k.caffemodel'
- net = caffe.Net(model_def, # defines the structure of the model
- model_weights, # contains the trained weights
- caffe.TEST) # use test mode (e.g., don't perform dropout)
- mu = np.array([104.00698793, 116.66876762, 122.67891434])
- print 'mean values:', zip('BGR', mu)
- # create transformer for the input called 'data'
- transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
- transformer.set_transpose('data', (2,0,1)) # move image channels to outermost dimension
- transformer.set_mean('data', mu) # subtract the dataset-mean value in each channel
- transformer.set_raw_scale('data', 255) # rescale from [0, 1] to [0, 255]
- transformer.set_channel_swap('data', (2,1,0)) # swap channels from RGB to BGR
- image = caffe.io.load_image('D:\\Development\\caffe-windows\\examples\\images\\cat_resized.jpg')
- transformed_image = transformer.preprocess('data', image)
- plt.imshow(image)
- # copy the image data into the memory allocated for the net
- net.blobs['data'].data[...] = transformed_image
- ### perform classification
- output = net.forward()
- score = output['score'][0] # the output probability vector for the first image in the batch
- classed = np.argmax(score, axis=0)
- names = dict()
- all_labels = open(net_root + '\\classes-59.txt').readlines()
- scores = np.unique(classed)
- labels = [all_labels[s] for s in scores]
- num_scores = len(scores)
- def rescore (c):
- """ rescore values from original score values (0-59) to values ranging from 0 to num_scores-1 """
- return np.where(scores == c)[0][0]
- rescore = np.vectorize(rescore)
- painted = rescore(classed)
- plt.figure(figsize=(10, 10))
- plt.imshow(painted, cmap=plt.cm.get_cmap('jet', num_scores))
- # setup legend
- formatter = plt.FuncFormatter(lambda val, loc: labels[val])
- plt.colorbar(ticks=range(0, num_scores), format=formatter)
- plt.clim(-0.5, num_scores - 0.5)
- plt.savefig('output.jpg')
- elapsed_time = time.time() - start_time
- print('Execution time: %.3f' % (elapsed_time))
- plt.imsave('../data/output.jpg', painted, cmap=plt.cm.get_cmap('jet', num_scores))
Add Comment
Please, Sign In to add comment