SinisterMJ

Caffe Segmentation Python

Sep 13th, 2016
87
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 2.84 KB | None | 0 0
  1. # set up Python environment: numpy for numerical routines, and matplotlib for plotting
  2. import time
  3. import numpy as np
  4. import matplotlib.pyplot as plt
  5. # display plots in this notebook
  6.  
  7. # set display defaults
  8. plt.rcParams['figure.figsize'] = (10, 10)        # large images
  9. plt.rcParams['image.interpolation'] = 'nearest'  # don't interpolate: show square pixels
  10. plt.rcParams['image.cmap'] = 'gray'  # use grayscale output rather than a (potentially misleading) color heatmap
  11.  
  12. import os
  13. import sys
  14. import caffe
  15.  
  16. caffe.set_mode_cpu()
  17. start_time = time.time()
  18. net_root = 'D:\\Development\\caffe-windows\\fcn.berkeleyvision.org-master\\voc-fcn8s'
  19.  
  20. model_def = net_root + '\\deploy.prototxt'
  21. model_weights = net_root + '\\fcn8s-heavy-40k.caffemodel'
  22.  
  23. net = caffe.Net(model_def,      # defines the structure of the model
  24.                 model_weights,  # contains the trained weights
  25.                 caffe.TEST)     # use test mode (e.g., don't perform dropout)
  26.  
  27. mu = np.array([104.00698793, 116.66876762, 122.67891434])
  28. print 'mean values:', zip('BGR', mu)
  29.  
  30. # create transformer for the input called 'data'
  31. transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
  32.  
  33. transformer.set_transpose('data', (2,0,1))  # move image channels to outermost dimension
  34. transformer.set_mean('data', mu)            # subtract the dataset-mean value in each channel
  35. transformer.set_raw_scale('data', 255)      # rescale from [0, 1] to [0, 255]
  36. transformer.set_channel_swap('data', (2,1,0))  # swap channels from RGB to BGR
  37.  
  38. image = caffe.io.load_image('D:\\Development\\caffe-windows\\examples\\images\\cat_resized.jpg')
  39.  
  40. transformed_image = transformer.preprocess('data', image)
  41. plt.imshow(image)
  42.  
  43. # copy the image data into the memory allocated for the net
  44. net.blobs['data'].data[...] = transformed_image
  45.  
  46. ### perform classification
  47. output = net.forward()
  48.  
  49. score = output['score'][0]  # the output probability vector for the first image in the batch
  50.  
  51. classed = np.argmax(score, axis=0)
  52.  
  53. names = dict()
  54. all_labels = open(net_root + '\\classes-59.txt').readlines()
  55. scores = np.unique(classed)
  56. labels = [all_labels[s] for s in scores]
  57. num_scores = len(scores)
  58.  
  59. def rescore (c):
  60.     """ rescore values from original score values (0-59) to values ranging from 0 to num_scores-1 """
  61.     return np.where(scores == c)[0][0]
  62.  
  63. rescore = np.vectorize(rescore)
  64.  
  65. painted = rescore(classed)
  66.  
  67. plt.figure(figsize=(10, 10))
  68. plt.imshow(painted, cmap=plt.cm.get_cmap('jet', num_scores))
  69.  
  70. # setup legend
  71. formatter = plt.FuncFormatter(lambda val, loc: labels[val])
  72. plt.colorbar(ticks=range(0, num_scores), format=formatter)
  73. plt.clim(-0.5, num_scores - 0.5)
  74. plt.savefig('output.jpg')
  75.  
  76. elapsed_time = time.time() - start_time
  77. print('Execution time: %.3f' % (elapsed_time))
  78. plt.imsave('../data/output.jpg', painted, cmap=plt.cm.get_cmap('jet', num_scores))
Add Comment
Please, Sign In to add comment