【FCN实践】04 预测

Fully Convolutional Networks for Semantic Segmentation

——————————————————————————————————————————

【FCN实践】01 常见问题 http://blog.csdn.net/binlearning/article/details/72854136

【FCN实践】02 模型迁移及初始化 http://blog.csdn.net/binlearning/article/details/72854244

【FCN实践】03 训练 http://blog.csdn.net/binlearning/article/details/72854407

【FCN实践】04 预测 http://blog.csdn.net/binlearning/article/details/72854583

【项目源码】https://github.com/binLearning/fcn_voc_32s

——————————————————————————————————————————

根据官方infer.py程序修改

import numpy as np
#from PIL import Image

import caffe

from scipy.misc import imread, imsave
from skimage.color import label2rgb

# Get the specified bit value
def bitget(byteval, idx):
return ((byteval & (1 << idx)) != 0)

# Create label-color map, label --- [R G B]
#  0 --- [  0   0   0],  1 --- [128   0   0],  2 --- [  0 128   0]
#  4 --- [128 128   0],  5 --- [  0   0 128],  6 --- [128   0 128]
#  7 --- [  0 128 128],  8 --- [128 128 128],  9 --- [ 64   0   0]
# 10 --- [192   0   0], 11 --- [ 64 128   0], 12 --- [192 128   0]
# 13 --- [ 64   0 128], 14 --- [192   0 128], 15 --- [ 64 128 128]
# 16 --- [192 128 128], 17 --- [  0  64   0], 18 --- [128  64   0]
# 19 --- [  0 192   0], 20 --- [128 192   0], 21 --- [  0  64 128]
def labelcolormap(N=256):
color_map = np.zeros((N, 3))
for n in xrange(N):
id_num = n
r, g, b = 0, 0, 0
for pos in xrange(8):
r = np.bitwise_or(r, (bitget(id_num, 0) << (7-pos)))
g = np.bitwise_or(g, (bitget(id_num, 1) << (7-pos)))
b = np.bitwise_or(b, (bitget(id_num, 2) << (7-pos)))
id_num = (id_num >> 3)
color_map[n, 0] = r
color_map[n, 1] = g
color_map[n, 2] = b
return color_map

def main():
# load image, switch to BGR, subtract mean, and make dims C x H x W for Caffe
im = Image.open('./data/test.jpg')
in_ = np.array(im, dtype=np.float32)
in_ = in_[:,:,::-1]
in_ -= np.array((104.00698793,116.66876762,122.67891434))
in_ = in_.transpose((2,0,1))

# load net
net = caffe.Net('./model/voc_fcn32s_deploy.prototxt',
'./snapshot/voc_fcn32s_iter_84980.caffemodel',
caffe.TEST)
# shape for input (data blob is N x C x H x W), set data
net.blobs['data'].reshape(1, *in_.shape)
net.blobs['data'].data[...] = in_
# run net and take argmax for prediction
net.forward()
out = net.blobs['score'].data[0].argmax(0).astype(np.uint8)

color_map = labelcolormap(21)

label_mask = label2rgb(out, colors=color_map[1:], bg_label=0)
label_mask[out == 0] = [0, 0, 0]
imsave('test_prediction.png', label_mask.astype(np.uint8))

if __name__ == '__main__':
main()

示例：