%createPointSeg Create PointSeg for semantic segmentation using deep learning.
%
%   pointseg is a convolutional neural network for semantic pointcloud
%   segmentation. It uses a pixelClassificationLayer to predict the
%   categorical label for every point in an input projected image.
%
%   Use pointsegLayers to create the network architecture for PointSeg. This
%   network must be trained using trainNetwork from Deep Learning Toolbox
%   before it can be used for semantic segmentation.
%
%   lgraph = pointsegLayers(imageSize, numClasses) returns U-Net layers
%   configured using the following inputs:
%
%   Inputs
%   ------
%   inputSize    - size of the network input image specified as a vector
%                  [H W] , where H and W are the image height and width,
%                  and C is the number of spherical image channels.
%
%   classNames   - names of classes the network should be configured to
%                  classify.
%
%   classWeights - weights of classes the network to balance the classes
%
% Notes
% -----
%
%
%   Example 1 - Create PointSeg with custom Input.
%   ------------------------------------------------------------
%   % Create pointSeg layers with an .
%   imageSize = [480 640 5];
%   classNames = [
%   "car"
%   "cyclist"
%   ];
%   classWeights = [1 1];
%   lgraph = pointsegLayers(imageSize, classNames, classWeights)
%
%   % Display network.
%   figure
%   plot(lgraph)
%

% References
% ----------
%
% [1] 

%   Copyright 2019 The MathWorks, Inc.

function lgraph = createPointSeg(inputSize, classNames, classWeights)

    numClasses = size(classNames,1);
    currentSize = inputSize(1:2);
    layers = [
        imageInputLayer(inputSize, 'Name', 'input')
        iConvBNreluLayers('conv1', 64, 3, 'Stride', [1 2])
        maxPooling2dLayer(3, 'Name', 'conv1_pool', 'Stride', [1 2], 'Padding', 'same')];
    lgraph = layerGraph(layers);

    lgraph = addLayers(lgraph, iConvBNreluLayers('conv_skip', 64, 1));
    lgraph = connectLayers(lgraph, 'input', 'conv_skip_conv');
    
    currentSize = currentSize./[1 4];
    
    [lgraph, fire1tail] = iAddFireModule('fire1', [16 64], ['conv1_pool'], lgraph);
    [lgraph, fire2tail] = iAddFireModule('fire2', [16 64], fire1tail, lgraph);
    [lgraph, SR1tail] = iAddSqueezeReweight('SR1', 128, currentSize, fire2tail, lgraph);
    
    lgraph = addLayers(lgraph, maxPooling2dLayer(3, 'Name', 'SR1_pool', 'Stride', [1 2], 'Padding', 'same'));
    lgraph = connectLayers(lgraph, SR1tail, 'SR1_pool');

    currentSize = currentSize./[1 2];

    [lgraph, fire3tail] = iAddFireModule('fire3', [32 128], 'SR1_pool', lgraph);
    [lgraph, fire4tail] = iAddFireModule('fire4', [32 128], fire3tail, lgraph);
    [lgraph, SR2tail] = iAddSqueezeReweight('SR2', 256, currentSize, fire4tail, lgraph);
    
    lgraph = addLayers(lgraph, maxPooling2dLayer(3, 'Name', 'SR2_pool', 'Stride', [1 2], 'Padding', 'same'));
    lgraph = connectLayers(lgraph, SR2tail, 'SR2_pool');

    currentSize = currentSize./[1 2];
    
    [lgraph, fire5tail] = iAddFireModule('fire5', [48 192], 'SR2_pool', lgraph);
    [lgraph, fire6tail] = iAddFireModule('fire6', [48 192], fire5tail, lgraph);
    [lgraph, fire7tail] = iAddFireModule('fire7', [64 256], fire6tail, lgraph);
    [lgraph, fire8tail] = iAddFireModule('fire8', [64 256], fire7tail, lgraph);
    
    [lgraph, SR3tail] = iAddSqueezeReweight('SR3', 512, currentSize, fire8tail, lgraph);
       
    [lgraph, ASPPtail] = iAddASPP('ASPP', 128, [1 6 9 12], currentSize, SR3tail, lgraph);
    [lgraph, ASPPfdtail] = iAddFireDeconvModule('fire_deconv_ASPP', [32 128], ASPPtail, lgraph);
    
    [lgraph, fd1tail] = iAddFireDeconvModule('fire_deconv1', [64 128], SR3tail, lgraph);
    [lgraph, add1tail] = iAddLayers({fd1tail, SR2tail}, lgraph);

    [lgraph, fuse1tail] = iConcatLayers({add1tail, ASPPfdtail}, lgraph);
    
    [lgraph, fd2tail] = iAddFireDeconvModule('fire_deconv2', [64 64], fuse1tail, lgraph);
    [lgraph, add2tail] = iAddLayers({fd2tail, SR1tail}, lgraph);

    [lgraph, fd3tail] = iAddFireDeconvModule('fire_deconv3', [16 32], add2tail, lgraph);
    [lgraph, add3tail] = iAddLayers({fd3tail, 'conv1_relu'}, lgraph);
    
    [lgraph, fd4tail] = iAddFireDeconvModule('fire_deconv4', [16 32], add3tail, lgraph);
    [lgraph, add4tail] = iAddLayers({fd4tail, 'conv_skip_relu'}, lgraph);
    
    layers = [
        dropoutLayer('Name', 'dropout')
        iConvBNreluLayers('convlast', numClasses, 3, 'BN', false)
        softmaxLayer('Name', 'softmax')
        pixelClassificationLayer('Name', 'pixellabels','Classes', classNames, 'ClassWeights', classWeights)];
    
    lgraph = addLayers(lgraph, [layers]);
    lgraph = connectLayers(lgraph, add4tail, 'dropout');
end


function [lgraph tail] = iConcatLayers(layerNames, lgraph)
%   Copyright 2019 The MathWorks, Inc.

    jName = [];
    pre = 'concat_';
    for i = 1:size(layerNames,2)
        jName = [jName cell2mat(layerNames(i))];
    end
    layer = depthConcatenationLayer(size(layerNames,2), 'Name', [pre jName]);
    lgraph = addLayers(lgraph,layer);
    
    for i = 1:size(layerNames,2)
        lgraph = connectLayers(lgraph, cell2mat(layerNames(i)), [pre jName '/in' num2str(i)]);
    end
    tail = [pre jName];
end


function [lgraph tail] = iAddLayers(layerNames, lgraph)
% IADDLAYERS joins all the layers in the input by an addtion layer
%

%   Copyright 2019 The MathWorks, Inc.

jName = [];
pre = 'fuse_';
for i = 1:size(layerNames,2)
    jName = [jName cell2mat(layerNames(i))];
end

layer = additionLayer(size(layerNames,2), 'Name', [pre jName]);
lgraph = addLayers(lgraph,layer);

for i = 1:size(layerNames,2)
    lgraph = connectLayers(lgraph, cell2mat(layerNames(i)), [pre jName '/in' num2str(i)]);
end
tail = [pre jName];
end


function layers = iConvBNreluLayers(name, numFilters, filterSize, varargin)
% iConvBNreluLayers   2D convolution layer for Convolutional Neural Networks
%
%   layers = iConvBNreluLayers(name, filterSize, numFilters) creates a layer
%   for 2D convolution. filterSize specifies the height and width of the
%   filters. It can be a scalar, in which case the filters will have the
%   same height and width, or a vector [h w] where h specifies the height
%   for the filters, and w specifies the width. numFilters specifies the
%   number of filters. name specifies the name prefix to the layers.
% 
%   layer = iConvBNreluLayers(name, filterSize, numFilters, 'PARAM1', VAL1, 'PARAM2', VAL2, ...) 
%   specifies optional parameter name/value pairs for creating the layer:
%
%       'Stride'                  - The step size for traversing the input
%                                   vertically and horizontally. This can
%                                   be a scalar, in which case the same
%                                   value is used for both dimensions, or
%                                   it can be a vector [u v] where u is the
%                                   vertical stride, and v is the
%                                   horizontal stride. The default is 
%                                   [1 1].
%       'DilationFactor'          - The up-sampling factor of the filter.
%                                   It corresponds to an effective filter
%                                   size of filterSize + (filterSize-1) .*
%                                   (dilationFactor-1). This can be a
%                                   scalar, in which case the same value is
%                                   used for both dimensions, or it can be
%                                   a vector [dHeight dWidth] where dHeight
%                                   is the vertical dilation, and dWidth is
%                                   the horizontal dilation. The default is
%                                   [1 1].
%        'Relu'                   - Add a Relu layer. This can be true or 
%                                   false. The Default is true
%        'BN'                     - Add a Batch Normalization layer. This 
%                                   can be true or false. The Default is 
%                                   true.
%
%
%   Example 1:
%       % Create a convolutional layer with 96 filters that have a height and
%       % width of 11, and use a stride of 4 in the horizontal and vertical 
%       % directions.
%
%       layer = iConvBNreluLayers('first' ,11, 96, 'Stride', 4);
%

%   Copyright 2019 The MathWorks, Inc.

p = inputParser;

defaultDialation = [1 1];
addOptional(p,'DilationFactor', defaultDialation);

defaultStride = [1 1];
addOptional(p,'Stride', defaultStride);

defaultRelu = true;
addOptional(p,'Relu', defaultRelu, @islogical);

defaultBN = true;
addOptional(p,'BN', defaultBN, @islogical);

p.KeepUnmatched = true;

parse(p,varargin{:})

if(p.Results.BN && p.Results.Relu)
    layers = [
        convolution2dLayer(filterSize, numFilters,'Padding','same','Name',[name, '_conv'],...
        'DilationFactor', p.Results.DilationFactor, 'Stride', p.Results.Stride)
        batchNormalizationLayer('Name',[name, '_BN'])
        reluLayer('Name',[name, '_relu'])];
elseif(~p.Results.BN && p.Results.Relu)
    layers = [
        convolution2dLayer(filterSize, numFilters,'Padding','same','Name',[name, '_conv'],...
        'DilationFactor', p.Results.DilationFactor, 'Stride', p.Results.Stride)
        reluLayer('Name',[name, '_relu'])];
elseif(p.Results.BN && ~p.Results.Relu)
    layers = [
        convolution2dLayer(filterSize, numFilters,'Padding','same','Name',[name, '_conv'],...
        'DilationFactor', p.Results.DilationFactor, 'Stride', p.Results.Stride)
        batchNormalizationLayer('Name',[name, '_BN'])];
else
    layers = [
        convolution2dLayer(filterSize, numFilters,'Padding','same','Name',[name, '_conv'],...
        'DilationFactor', p.Results.DilationFactor, 'Stride', p.Results.Stride)];
end

end


function [lgraph tail] = iAddASPP(name, numFilters, dilations, inSize, inputLayerName, lgraph)
% An ASPP module can be visualized as follows:
%
%                                |
%              +--------+--------+-------+-------+
%              |        |        |       |       |
%              |        |        |       |       |
%              v        v        v       v       v
%           +--+---+ +--+---+ +--+--+ +--+--+ +--+---+
%           |ASPP  | |ASPP  | |ASPP | |ASPP | |ASPP  |
%           |Branch| |Branch| |Branch |Branch |Branch|
%           |D1    | |D6    | |D12  | |D18  | |GAP   |
%           +--+---+ +--+---+ +--+--+ +--+--+ +---+--+
%              |        |        |       |        |
%              v        v        v       v        v
%           +-----+-----+--------+-------+--------+--+
%           |          Depth concatenation           |
%           +--------------------+-------------------+
%                                |   
%                                v
%                         +------+------+
%                         |  conv 1x1   |
%                         +------+------+
%
%
%   Copyright 2015-2018 The MathWorks, Inc.

for dno = 1:size(dilations,2)
    d = dilations(dno);
    layer = iConvBNreluLayers([name '_' num2str(d)], numFilters, 3,  'DilationFactor' ,[d d]);
    lgraph = addLayers(lgraph,layer);
    lgraph = connectLayers(lgraph, inputLayerName, [name '_' num2str(d) '_conv']);
end

layers = [averagePooling2dLayer(inSize, 'Name', [name '_GAP'])
    iConvBNreluLayers([name '_GAP'], numFilters, 1)
    resize2dLayer('OutputSize', inSize, 'Name', [name '_GAP_resize']);
    depthConcatenationLayer(5, 'Name', [name, '_concat'])
    iConvBNreluLayers([name '_conv'], numFilters, 1, 'BN', false)];

lgraph = addLayers(lgraph,layers);
lgraph = connectLayers(lgraph, inputLayerName, [name '_GAP']);

for dno = 1:size(dilations,2)
    d = dilations(dno);
    lgraph = connectLayers(lgraph, [name '_' num2str(d) '_relu'], [name '_concat/in' num2str(dno+1)]);
end
 tail = [name '_conv_relu'];
end


function [lgraph tail]=  iAddFireModule(name, numFilters, inputLayerName, lgraph)
% IADDFIREMODULE adds a fire module to the layer graph
%
% An fire module can be visualized as follows:
%
%                     |
%                +----+----+
%                |Conv 1x1 |
%                +----+----+
%                     |
%              +------+------+
%              |             |
%              v             v
%         +----+----+   +----+----+
%         |Conv 1x1 |   |Conv 3x3 |
%         +----+----+   +----+----+
%              |             |    
%              v             v    
%           +-----+-----+--------+
%           | Depth concatenation|
%           +--------------------+
%                     |
%                     v
%

%   Copyright 2019 The MathWorks, Inc.

layers = [
    iConvBNreluLayers([name,'_squeeze1x1'], numFilters(1), 1)
    iConvBNreluLayers([name,'_expand1x1'], numFilters(2), 1)
    depthConcatenationLayer(2, 'Name', [name, '_concat'])];
lgraph = addLayers(lgraph,layers);


layer = iConvBNreluLayers([name,'_expand3x3'], numFilters(2), 3);
lgraph = addLayers(lgraph,layer);

lgraph = connectLayers(lgraph, inputLayerName, [name,'_squeeze1x1_conv']);
lgraph = connectLayers(lgraph,[name,'_squeeze1x1_relu'],[name,'_expand3x3_conv']);
lgraph = connectLayers(lgraph,[name,'_expand3x3_relu'],[name, '_concat/in2']);
tail = [name, '_concat'];
end


function [lgraph tail] =  iAddFireDeconvModule(name, numFilters, inputLayerName, lgraph)
% IADDFIREDECONVMODULE adds a fire module to the layer graph
%
% An fire module can be visualized as follows:
%
%                     |
%                +----+----+
%                |Conv 1x1 |
%                +----+----+
%                     |
%                +----+----+
%                |Deconv x2|
%                +----+----+
%                     |
%              +------+------+
%              |             |
%              v             v
%         +----+----+   +----+----+
%         |Conv 1x1 |   |Conv 3x3 |
%         +----+----+   +----+----+
%              |             |
%              v             v
%           +-----+-----+--------+
%           | Depth concatenation|
%           +--------------------+
%                     |
%                     v
%

%   Copyright 2019 The MathWorks, Inc.
layers = [
    iConvBNreluLayers([name,'_squeeze1x1'], numFilters(1), 1)
    transposedConv2dLayer([1 4], numFilters(1), 'Cropping', 'same', ...
    'Stride', [1 2], 'Name', [name, '_deconv'])
    reluLayer('Name', [name, '_deconv_relu'])
    iConvBNreluLayers([name,'_expand1x1'], numFilters(2), 1)
    depthConcatenationLayer(2, 'Name', [name, '_concat'])];
lgraph = addLayers(lgraph,layers);


layer = iConvBNreluLayers([name,'_expand3x3'], numFilters(2), 3);
lgraph = addLayers(lgraph,layer);

lgraph = connectLayers(lgraph, inputLayerName, [name,'_squeeze1x1_conv']);
lgraph = connectLayers(lgraph, [name,'_deconv_relu'], [name,'_expand3x3_conv']);
lgraph = connectLayers(lgraph, [name,'_expand3x3_relu'], [name, '_concat/in2']);
tail = [name, '_concat'];
end


function [lgraph tail] = iAddSqueezeReweight(name, numFilters, inSize, inputLayerName, lgraph)
%   Copyright 2019 The MathWorks, Inc.

layers = [
    averagePooling2dLayer(inSize,'Name', [name '_GAP'])
    fullyConnectedLayer(numFilters/2, 'Name', [name '_FC_1'])
    sigmoidLayer('Name',[name '_SIG_1'])
    fullyConnectedLayer(numFilters, 'Name', [name '_FC_2'])
    sigmoidLayer('Name',[name '_SIG_2'])
    multiplicationLayer(2,'Name',[name '_REW'])];
tail =  [name '_REW'];

lgraph = addLayers(lgraph,layers);

lgraph = connectLayers(lgraph, inputLayerName, [name,'_GAP']);
lgraph = connectLayers(lgraph, inputLayerName, [name, '_REW/in2']);
end