Permalink
Switch branches/tags
Nothing to show
Find file Copy path
Fetching contributors…
Cannot retrieve contributors at this time
98 lines (77 sloc) 3.02 KB
% This script demos how to use the pre-trained models to
% obtain the predicted segmentations
close all; clc; clear;
addpath(genpath('visualizationCode'));
% path to caffe (compile matcaffe first, or you could use python wrapper instead)
addpath 'yourcaffe/matlab'
% select the pre-trained model. Use 'FCN' for
% the Fully Convolutional Network or 'Dilated' for DilatedNet
% You can download the FCN model at
% http://sceneparsing.csail.mit.edu/model/caffe/FCN_iter_160000.caffemodel
% and the DilatedNet model at
% http://sceneparsing.csail.mit.edu/model/caffe/DilatedNet_iter_120000.caffemodel
model_type = 'FCN'; %Dilated'
if (strcmp(model_type, 'FCN'))
model_definition = 'models/deploy_FCN.prototxt';
model_weights = 'FCN_iter_160000.caffemodel';
elseif (strcmp(model_type, 'Dilated'))
model_definition = 'models/deploy_DilatedNet.prototxt';
model_weights = 'DilatedNet_iter_120000.caffemodel';
end
disp(model_definition)
prediction_folder = sprintf('predictions_%s', model_type);
% initialize the network
net = caffe.Net(model_definition, model_weights, 'test');
% path to image(.jpg) and annotation(.png) and generated prediction(.png)
pathImg = fullfile('sampleData', 'images');
pathAnno = fullfile('sampleData', 'annotations');
pathPred = fullfile('sampleData', prediction_folder);
if (~exist(pathPred, 'dir'))
mkdir(pathPred);
end
% load class names
load('objectName150.mat');
% load pre-defined colors
load('color150.mat');
filesImg = dir(fullfile(pathImg, '*.jpg'));
for i = 1: numel(filesImg)
% read image
fileImg = fullfile(pathImg, filesImg(i).name);
fileAnno = fullfile(pathAnno, strrep(filesImg(i).name, '.jpg', '.png'));
filePred = fullfile(pathPred, strrep(filesImg(i).name, '.jpg', '.png'));
im = imread(fileImg);
imAnno = imread(fileAnno);
% resize image to fit model description
im_inp = double(imresize(im, [384,384]));
% change RGB to BGR
im_inp = im_inp(:,:,end:-1:1);
% substract mean and transpose
im_inp = cat(3, im_inp(:,:,1)-109.5388, im_inp(:,:,2)-118.6897, im_inp(:,:,3)-124.6901);
im_inp = permute(im_inp, [2,1,3]);
% obtain predicted image and resize to original size
imPred = net.forward({im_inp});
[~, imPred] = max(imPred{1},[],3);
imPred = uint8(imPred')-1;
imPred = imresize(imPred, [size(im,1), size(im,2)], 'nearest');
imwrite(imPred, filePred);
% color encoding
rgbPred = colorEncode(imPred, colors);
rgbAnno = colorEncode(imAnno, colors);
% colormaps
colormap = colorMap(imPred, imAnno, objectNames);
% plot
set(gcf, 'Name', [fileImg ' [Press any key to the next image...]'], 'NumberTitle','off');
subplot(231);
imshow(im); title('Image');
subplot(232);
imshow(imPred); title('Prediction-gray');
subplot(233);
imshow(imAnno); title('Annotation-gray');
subplot(234);
imshow(colormap); title('Colormap');
subplot(235);
imshow(rgbPred); title('Prediction-color');
subplot(236);
imshow(rgbAnno); title('Annotation-color');
waitforbuttonpress;
end