In [ ]:
import os
os.environ["PYTHONPATH"] = '/opt/SegNet/caffe-segnet/python'
!sudo unlink /usr/local/cuda
!sudo ln -s /usr/local/cuda-6.5 /usr/local/cuda
In [ ]:
Type = 'segnet' # Options: segnet, segnetB, bayesian, bayesianB, webdemo
NUM_IMAGES = 255
MODEL_PATH = '/opt/SegNet/Models/'
WEIGHT_PATH = '/opt/SegNet/Models/Inference/'
if Type == 'segnet':
model = MODEL_PATH + 'segnet_inference.prototxt'
weights = WEIGHT_PATH + 'test_weights_segnet.caffemodel'
elif Type == 'segnetB':
model = MODEL_PATH + 'segnet_basic_inference.prototxt'
weights = WEIGHT_PATH + 'test_weights_segnet_basic.caffemodel'
elif Type == 'bayesian':
model = MODEL_PATH + 'bayesian_segnet_inference.prototxt'
weights = WEIGHT_PATH + 'test_weights_bayesian.caffemodel'
elif Type == 'bayesianB':
model = MODEL_PATH + 'bayesian_segnet_basic_inference.prototxt'
weights = WEIGHT_PATH + 'test_weights_bayesian_basic.caffemodel'
elif Type == 'webdemo':
model = MODEL_PATH + 'webdemo_inference.prototxt'
weights = WEIGHT_PATH + 'test_weights_webdemo.caffemodel'
if Type == 'segnet' or Type == 'segnetB':
iterations = NUM_IMAGES
if Type == 'bayesian' or Type == 'bayesianB':
colours = MODEL_PATH + '../Scripts/camvid11.png'
data = MODEL_PATH + '../CamVid/test.txt'
In [ ]:
import numpy as np
import matplotlib
import matplotlib.colors as colors
import matplotlib.cm as cmx
import matplotlib.pyplot as plt
import os.path, scipy, json, math, pylab, sys, cv2
import scipy.io as sio
from sklearn.preprocessing import normalize
caffe_root = '/opt/SegNet/caffe-segnet/'
sys.path.insert(0, caffe_root + 'python')
import caffe
%matplotlib inline
In [ ]:
caffe.set_mode_gpu()
net = caffe.Net(model, weights, caffe.TEST)
if Type == 'bayesian' or Type == 'bayesianB':
input_shape = net.blobs['data'].data.shape
label_colours = cv2.imread(colours).astype(np.uint8)
In [ ]:
# For SegNet
if Type == 'segnet' or Type == 'segnetB':
print ('Performing SegNet Classification')
for i in range(0, iterations):
net.forward()
image = net.blobs['data'].data
label = net.blobs['label'].data
predicted = net.blobs['prob'].data
image = np.squeeze(image[0,:,:,:])
output = np.squeeze(predicted[0,:,:,:])
ind = np.argmax(output, axis=0)
r = ind.copy()
g = ind.copy()
b = ind.copy()
r_gt = label.copy()
g_gt = label.copy()
b_gt = label.copy()
Sky = [128,128,128]
Building = [128,0,0]
Pole = [192,192,128]
Road_marking = [255,69,0]
Road = [128,64,128]
Pavement = [60,40,222]
Tree = [128,128,0]
SignSymbol = [192,128,128]
Fence = [64,64,128]
Car = [64,0,128]
Pedestrian = [64,64,0]
Bicyclist = [0,128,192]
Unlabelled = [0,0,0]
label_colours = np.array([Sky, Building, Pole, Road, Pavement,
Tree, SignSymbol, Fence, Car, Pedestrian,
Bicyclist, Unlabelled])
for l in range(0,11):
r[ind==l] = label_colours[l,0]
g[ind==l] = label_colours[l,1]
b[ind==l] = label_colours[l,2]
r_gt[label==l] = label_colours[l,0]
g_gt[label==l] = label_colours[l,1]
b_gt[label==l] = label_colours[l,2]
rgb = np.zeros((ind.shape[0], ind.shape[1], 3))
rgb[:,:,0] = r/255.0
rgb[:,:,1] = g/255.0
rgb[:,:,2] = b/255.0
rgb_gt = np.zeros((ind.shape[0], ind.shape[1], 3))
rgb_gt[:,:,0] = r_gt/255.0
rgb_gt[:,:,1] = g_gt/255.0
rgb_gt[:,:,2] = b_gt/255.0
image = image/255.0
image = np.transpose(image, (1,2,0))
output = np.transpose(output, (1,2,0))
image = image[:,:,(2,1,0)]
# scipy.misc.toimage(rgb, cmin=0.0,
# cmax=255).save(IMAGE_FILE+'_segnet.png')
plt.figure()
plt.imshow(image, vmin=0, vmax=1)
plt.figure()
plt.imshow(rgb_gt, vmin=0, vmax=1)
plt.figure()
plt.imshow(rgb, vmin=0, vmax=1)
plt.axis('off')
plt.show()
print 'Success!'
# For Bayesian SegNet
elif Type == 'bayesian' or Type == 'bayesianB':
print ('Performing Bayesian SegNet Classification')
with open(data) as f:
for line in f:
input_image_file, ground_truth_file = line.split()
input_image_raw = caffe.io.load_image(input_image_file)
ground_truth = cv2.imread(ground_truth_file, 0)
input_image = caffe.io.resize_image(input_image_raw, (input_shape[2], input_shape[3]))
input_image = input_image*255
input_image = input_image.transpose((2,0,1))
input_image = input_image[(2,1,0),:,:]
input_image = np.asarray([input_image])
input_image = np.repeat(input_image, input_shape[0], axis=0)
# out = net.forward_all(data=input_image)
out = net.forward()
predicted = net.blobs['prob'].data
output = np.mean(predicted,axis=0)
uncertainty = np.var(predicted,axis=0)
ind = np.argmax(output, axis=0)
segmentation_ind_3ch = np.resize(ind,(3,input_shape[2],input_shape[3]))
segmentation_ind_3ch = segmentation_ind_3ch.transpose(1,2,0).astype(np.uint8)
segmentation_rgb = np.zeros(segmentation_ind_3ch.shape, dtype=np.uint8)
gt_ind_3ch = np.resize(ground_truth,(3,input_shape[2],input_shape[3]))
gt_ind_3ch = gt_ind_3ch.transpose(1,2,0).astype(np.uint8)
gt_rgb = np.zeros(gt_ind_3ch.shape, dtype=np.uint8)
cv2.LUT(segmentation_ind_3ch,label_colours,segmentation_rgb)
cv2.LUT(gt_ind_3ch,label_colours,gt_rgb)
uncertainty = np.transpose(uncertainty, (1,2,0))
average_unc = np.mean(uncertainty,axis=2)
min_average_unc = np.min(average_unc)
max_average_unc = np.max(average_unc)
max_unc = np.max(uncertainty)
plt.imshow(input_image_raw,vmin=0, vmax=255)
plt.figure()
plt.imshow(segmentation_rgb,vmin=0, vmax=255)
plt.figure()
plt.imshow(gt_rgb,vmin=0, vmax=255)
plt.set_cmap('bone_r')
plt.figure()
plt.imshow(average_unc,vmin=0, vmax=max_average_unc)
plt.show()
# uncomment to save results
#scipy.misc.toimage(segmentation_rgb, cmin=0.0, cmax=255.0).save(IMAGE_FILE+'_segnet_segmentation.png')
#cm = matplotlib.pyplot.get_cmap('bone_r')
#matplotlib.image.imsave(input_image_file+'_segnet_uncertainty.png',average_unc,cmap=cm, vmin=0, vmax=max_average_unc)
print 'Processed: ', input_image_file
print 'Success!'
# For webdemo and other test cases
elif type == 'webdemo':
print 'Success'