%% Code Generation for Semantic Segmentation Application on Intel CPUs That Uses U-Net
% 
% This example demonstrates code generation for an image segmentation application 
% that uses deep learning. It uses the |codegen| command to generate a MEX 
% function that performs prediction by using the deep learning
% network U-Net for image segmentation.
%
% For a similar example that demonstrates segmentation of images by using
% U-Net but does not use the |codegen| command, see 
% <matlab:web(fullfile(docroot,'images/multispectral-semantic-segmentation-using-deep-learning.html'))
% Semantic Segmentation of Multispectral Images Using Deep Learning>.
%
% Copyright 2019-2020 The MathWorks, Inc.
%% Prerequisites
% * Xeon processor with support for Intel Advanced Vector Extensions 2
% (Intel AVX2) instructions
% * MATLAB&reg; Coder&trade;, for C++ code generation
% * MATLAB Coder(TM) Interface for Deep Learning Libraries support package
% * Deep Learning Toolbox&trade;, for using the |DAGNetwork| object
%
% This example is supported on Linux(R) and Windows(R) platforms.
%
% This example uses the Intel MKL-DNN library that ships with MATLAB and
% generates a MEX function for semantic segmentation.
% 
% This example is not supported in MATLAB Online.

%% Overview of U-Net
% 
% U-Net [1] is a type of convolutional neural network (CNN)that is
% designed for semantic image segmentation. In U-Net, the initial series of
% convolutional layers are interspersed with max pooling layers,
% successively decreasing the resolution of the input image. These layers
% are followed by a series of convolutional layers interspersed with
% upsampling operators, successively increasing the resolution of the input
% image. The combination of these two series paths forms a U-shaped graph. The 
% network was originally trained to perform
% prediction for biomedical image segmentation applications. This
% example demonstrates the ability of the network to track changes in
% forest cover over time. Environmental agencies track deforestation 
% to assess and qualify the environmental and ecological health of a region.
% 
% Deep-learning-based semantic segmentation can yield a precise measurement of 
% vegetation cover from high-resolution aerial photographs. One of the challenges is 
% differentiating classes that have similar visual characteristics, such as trying 
% to classify a green pixel as grass, shrubbery, or tree. To increase classification 
% accuracy, some data sets contain multispectral images that provide additional 
% information about each pixel. For example, the Hamlin Beach State Park data 
% set supplements the color images with near-infrared channels that provide a 
% clearer separation of the classes.
%
% This example uses the Hamlin Beach State Park Data [2] along with a
% pretrained U-Net network in order to correctly classify each pixel.
%
% The U-Net this example uses is trained to segment pixels belonging to 18 classes which 
% includes: 
%
%  0. Other Class/Image Border      7. Picnic Table         14. Grass
%  1. Road Markings                 8. Black Wood Panel     15. Sand
%  2. Tree                          9. White Wood Panel     16. Water (Lake)
%  3. Building                     10. Orange Landing Pad   17. Water (Pond)
%  4. Vehicle (Car, Truck, or Bus) 11. Water Buoy           18. Asphalt (Parking Lot/Walkway)
%  5. Person                       12. Rocks
%  6. Lifeguard Chair              13. Other Vegetation

%% Get Pretrained U-Net DAG Network Object
%
trainedUnet_url = 'https://www.mathworks.com/supportfiles/vision/data/multispectralUnet.mat';
downloadTrainedUnet(trainedUnet_url,pwd);

%%
%
ld = load("trainedUnet/multispectralUnet.mat");
net = ld.net;

%%
% The DAG network contains 58 layers including convolution, max pooling, depth 
% concatenation, and pixel classification output layers. To display an 
% interactive visualization of the deep learning network architecture, use the 
% <docid:nnet_ref#mw_8d52b67d-b6b3-4c62-b573-56fdf4dce6a0 analyzeNetwork> 
% function.

%   analyzeNetwork(net);

%% The |segmentImageUnet| Entry-Point Function
%
% The
% <matlab:edit(fullfile(matlabroot,'examples','deeplearning_shared','main','segmentImageUnet.m'))
% segmentImageUnet.m> entry-point function performs semantic segmentation on the input
% image for each patch of a fixed size by using the multispectralUnet network contained in the
% |multispectralUnet.mat| file. This function loads the network object from
% the |multispectralUnet.mat| file into a persistent variable |mynet|. The function 
% reuses this persistent variable in subsequent prediction calls.
%
type('segmentImageUnet.m')

%% Prepare Data
%
% Download the Hamlin Beach State Park data.
if ~exist(fullfile(pwd,'data'),'dir')
    url = 'http://www.cis.rit.edu/~rmk6217/rit18_data.mat';
    downloadHamlinBeachMSIData(url,pwd+"/data/");
end

%%
%
% Load and examine the data in MATLAB.
load(fullfile(pwd,'data','rit18_data','rit18_data.mat'));

% Examine data
whos test_data

%%
% The image has seven channels. The RGB color channels are
% the fourth, fifth, and sixth image channels. The first three channels
% correspond to the near-infrared bands and highlight different components
% of the image based on their heat signatures. Channel 7 is a mask that
% indicates the valid segmentation region.
%
% The multispectral image data is arranged as
% numChannels-by-width-by-height arrays. In MATLAB, multichannel
% images are arranged as width-by-height-by-numChannels arrays. To reshape
% the data so that the channels are in the third dimension, use the helper
% function, |switchChannelsToThirdPlane|.
test_data  = switchChannelsToThirdPlane(test_data);

% Confirm data has the correct structure (channels last).
whos test_data

%% Generate MEX
%
% To generate a MEX function for the
% <matlab:edit(fullfile(matlabroot,'examples','deeplearning_shared','main','segmentImageUnet.m')) segmentImageUnet.m> 
% entry-point function, create a code configuration object |cfg| for MEX code generation. 
% Set the target language to C++. Use the
% <docid:gpucoder_ref#mw_e8e85f8e-8dde-45b6-9ec5-f121a79dc48f coder.DeepLearningConfig>
% function to create an MKL-DNN deep learning configuration object and
% assign it to the |DeepLearningConfig| property of |cfg|. Run the |codegen| command specifying an input size of
% [12446,7654,7] and a patch size of [1024,1024]. These
% values correspond to the size of the entire |test_data| variable. The smaller patch sizes
% speed up inference. To see how the patches are calculated, see the 
% |segmentImageUnet| entry-point function.
cfg = coder.config('mex');
cfg.TargetLang = 'C++';
cfg.DeepLearningConfig = coder.DeepLearningConfig('mkldnn');
codegen -config cfg segmentImageUnet -args {ones(size(test_data),'uint16'),coder.Constant([1024 1024])} -report

%% Run Generated MEX to Predict Results for |test_data|
% The |segmentImageUnet| function accepts |test_data|
% and a vector containing the dimensions of the patch size as inputs. The 
% function divides the image into patches, predicts the pixels in a
% particular patch, and finally combines all the patches. 
% Because of the large size of |test_data| 
% (12446x7654x7), it is easier to process the image in patches.
segmentedImage = segmentImageUnet_mex(test_data,[1024 1024]);

%%
% To extract only the valid portion of the segmentation, multiply the
% segmented image by the mask channel of the test data.
segmentedImage = uint8(test_data(:,:,7)~=0) .* segmentedImage;

%%
% Remove the noise and stray pixels by using the |medfilt2| function.
segmentedImage = medfilt2(segmentedImage,[5,5]);

%% Display U-Net Segmented test_data 
% This line of code creates a vector of the class names:
classNames = net.Layers(end).Classes; 

%%
% Overlay the labels on the segmented RGB test image and add a color bar to the
% segmentation image.

% Display input data

figure(1);
imshow(histeq(test_data(:,:,1:3)));
title('Input Image');
cmap = jet(numel(classNames));
segmentedImageOut = labeloverlay(imadjust(test_data(:,:,4:6),[0 0.6],[0.1 0.9],0.55),segmentedImage,'Transparency',0,'Colormap',cmap);

% Display segmented data

figure(2);
imshow(segmentedImageOut);
title('Segmented Image Output');
N = numel(classNames);
ticks = 1/(N*2):1/N:1;
colorbar('TickLabels',cellstr(classNames),'Ticks',ticks,'TickLength',0,'TickLabelInterpreter','none');
colormap(cmap)
title('Segmented Image using Mkldnn');
segmentedImageOverlay = labeloverlay(imadjust(test_data(:,:,4:6),[0 0.6],[0.1 0.9],0.55),segmentedImage,'Transparency',0.7,'Colormap',cmap);
figure(3);
imshow(segmentedImageOverlay);
title('Segmented Overlay Image');

%% References
% [1] Ronneberger, Olaf, Philipp Fischer, and Thomas Brox. "U-Net:
% Convolutional Networks for Biomedical Image Segmentation." _arXiv
% preprint arXiv:1505.04597,_ 2015.
%
% [2] Kemker, R., C. Salvaggio, and C. Kanan. "High-Resolution
% Multispectral Dataset for Semantic Segmentation." CoRR, abs/1703.01918,
% 2017.