# -*- coding: utf-8 -*
import numpy as np
import matplotlib.pyplot as plt
import os.path
import scipy
import argparse
import math
import cv2
import sys
import timesys.path.append('/usr/local/lib/python2.7/site-packages')
# Make sure that caffe is on the python path:
caffe_root = '/media/lbk/娱乐/seg-env/caffe-segnet/'
sys.path.insert(0, caffe_root + 'python')
import caffe# Import arguments
#deploy='Example_Models/segnet_model_driving_webdemo.prototxt'
#weights='Example_Moudels/segnet_weights_driving_webdemo.caffemodel'
#colours='Scripts/camvid12.png'
#net = caffe.Net(deploy,weights,caffe.TEST)# Import arguments
parser = argparse.ArgumentParser()
parser.add_argument('--model', type=str, required=True)
parser.add_argument('--weights', type=str, required=True)
parser.add_argument('--colours', type=str, required=True)
args = parser.parse_args()net = caffe.Net(args.model,
args.weights,
caffe.TEST)
#caffe.set_mode_gpu()input_shape = net.blobs['data'].data.shape
output_shape = net.blobs['argmax'].data.shapelabel_colours = cv2.imread(args.colours).astype(np.uint8)#cv2.namedWindow("Input")
#cv2.namedWindow("SegNet")cap = cv2.VideoCapture(0) # Change this to your webcam ID, or file name for your video filerval = Trueframe = cv2.imread('/media/lbk/娱乐/seg-env/caffe-segnet/segnet/Example_Models/123.png')
frame = cv2.resize(frame, (input_shape[3],input_shape[2]))
input_image = frame.transpose((2,0,1))
# input_image = input_image[(2,1,0),:,:] # May be required, if you do not open your data with opencv
input_image = np.asarray([input_image])
out = net.forward_all(data=input_image)segmentation_ind = np.squeeze(net.blobs['argmax'].data)
segmentation_ind_3ch = np.resize(segmentation_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)cv2.LUT(segmentation_ind_3ch,label_colours,segmentation_rgb)
#这里应该变成小数存储了，看来opencv对于小数也是热图显示，但是保存还是黑白的图
segmentation_rgb = segmentation_rgb.astype(float)/255#cv2.imwrite('output.jpg',segmentation_rgb)
#cv2.imshow("Input", frame)
#cv2.imshow("SegNet", segmentation_rgb)
#cv2.imwrite('output.jpg',segmentation_rgb)
#这里使用plt显示与保存，比cv2好点，并且不会出现进程卡住的情况
plt.imshow(segmentation_rgb)
plt.savefig('output.png')
plt.show()