%% Lidar Point Cloud Semantic Segmentation Using PointSeg Deep Learning Network
% This example shows how to train a PointSeg semantic segmentation network on 3-D organized lidar point cloud data.
% PointSeg [1] is a convolutional neural network (CNN) for performing end-to-end semantic segmentation of road objects based on an organized lidar point cloud. By using methods such as atrous spatial pyramid pooling (ASPP) and squeeze-and-excitation blocks, the network provides improved segmentation results. The training procedure shown in this example requires 2-D spherical projected images as inputs to the deep learning network.
% This example uses a highway scene data set collected using an Ouster OS1 sensor. It contains organized lidar point cloud scans of highway scenes and corresponding ground truth labels for car and truck objects. The size of the data file is approximately 760 MB.
%
%% Download Lidar Data Set
% Execute this code to download the highway scene data set. The data set contains 1617 point clouds stored as pointCloud objects in a cell array. Corresponding ground truth data, which is attached to the example, contains bounding box information of cars and trucks in each point cloud.

url = 'https://www.mathworks.com/supportfiles/lidar/data/WPI_LidarData.tar.gz';

outputFolder = fullfile(tempdir,'WPI');
lidarDataTarFile = fullfile(outputFolder,'WPI_LidarData.tar.gz');

if ~exist(lidarDataTarFile, 'file') 
    mkdir(outputFolder);
    
    disp('Downloading WPI Lidar driving data (760 MB)...');
    websave(lidarDataTarFile, url);
    untar(lidarDataTarFile,outputFolder); 
end

% Check if tar.gz file is downloaded, but not uncompressed.
if ~exist(fullfile(outputFolder, 'WPI_LidarData.mat'), 'file')
    untar(lidarDataTarFile,outputFolder);
end
lidarData = load(fullfile(outputFolder, 'WPI_LidarData.mat'));

groundTruthData = load('WPI_LidarGroundTruth.mat');

% Note: Depending on your Internet connection, the download process can take some time. The code suspends MATLAB® execution until the download process is complete. Alternatively, you can download the data set to your local disk using your web browser, and then extract WPI_LidarData. To use the file you downloaded from the web, change the outputFolder variable in the code to the location of the downloaded file.
%
%% Download Pretrained Network
% Download the pretrained network to avoid having to wait for training to complete. If you want to train the network, set the doTraining variable to true.

doTraining = false;
if ~doTraining && ~exist('trainedPointSegNet.mat','file')
    disp('Downloading pretrained network (14 MB)...');
    pretrainedURL = 'https://www.mathworks.com/supportfiles/lidar/data/trainedPointSegNet.mat';
    websave('trainedPointSegNet.mat', pretrainedURL);
end

%
%% Prepare Data for Training
%% Load Lidar Point Clouds and Class Labels
% Use the helperGenerateTrainingData supporting function, attached to this example, to generate training data from the lidar point clouds. The function uses point cloud and bounding box data to create five-channel input images and pixel label images. To create the pixel label images, the function selects points inside the bounding box and labels them with the bounding box class ID.  Each training image is specified as a 64-by-1024-by-5 array:
% The height of each image is 64 pixels.
% The width of each image is 1024 pixels.
% Each image has 5 channels. The five channels specify the 3-D coordinates of the point cloud, intensity, and range: r=\sqrt{x^{2\;} +y^2 +z^2 }.
% A visual representation of the training data follows.
%
% Generate the five-channel training images and pixel label images.

imagesFolder = fullfile(outputFolder, 'images');
labelsFolder = fullfile(outputFolder, 'labels');

helperGenerateTrainingData(lidarData, groundTruthData, imagesFolder, labelsFolder); 

% The five-channel images are saved as MAT files. Pixel labels are saved as PNG files.
% Note: Processing can take some time. The code suspends MATLAB® execution until processing is complete.
%
%% Create ImageDatastore and PixelLabelDatastore
% Use the ImageDatastore object to extract and store the five channels of the 2-D spherical images using the helperImageMatReader supporting function, which is a custom MAT file reader. This function is attached to this example as a supporting file.

imds = imageDatastore(imagesFolder, ...
         'FileExtensions', '.mat', ...
         'ReadFcn', @helperImageMatReader);

%
% Use the PixelLabelDatastore object to store pixel-wise labels from the label images. The object maps each pixel label to a class name. In this example, cars and trucks are the only objects of interest; all other pixels are the background. Specify these classes (car, truck, and background) and assign a unique label ID to each class.

classNames = [
    "background"
    "car"
    "truck"
];

numClasses = numel(classNames);

% Specify label IDs from 1 to the number of classes.
labelIDs = 1 : numClasses;

pxds = pixelLabelDatastore(labelsFolder, classNames, labelIDs);

%
% Load and display one of the labeled images by overlaying it on the corresponding intensity image using the helperDisplayLidarOverlayImage function, defined in the Supporting Functions section of this example.

imageNumber = 225;

% Point cloud (channels 1, 2, and 3 are for location, channel 4 is for intensity).
I = readimage(imds, imageNumber);

labelMap = readimage(pxds, imageNumber);
figure;
helperDisplayLidarOverlayImage(I, labelMap, classNames);
title('Ground Truth');

%
%% Prepare Training, Validation, and Test Sets
% Use the helperPartitionLidarData supporting function, attached to this example, to split the data into training, validation, and test sets that contain 970, 216, and 431 images, respectively.

[imdsTrain, imdsVal, imdsTest, pxdsTrain, pxdsVal, pxdsTest] = ...
    helperPartitionLidarData(imds, pxds);

% Use the combine function to combine the pixel and image datastores for the training and validation data sets.

trainingData = combine(imdsTrain, pxdsTrain); 
validationData = combine(imdsVal, pxdsVal);

%
%% Data Augmentation
% Data augmentation is used to improve network accuracy by randomly transforming the original data during training. By using data augmentation, you can add more variety to the training data without actually having to increase the number of labeled training samples.
% Augment the training data using the transform function with custom preprocessing operations specified by the augmentData function, defined in the Supporting Functions section of this example. This function randomly flips the spherical 2-D image and associated labels in the horizontal direction. Apply data augmentation to only the training data set.

augmentedTrainingData = transform(trainingData, @(x) augmentData(x));

%
%% Balance Classes Using Class Weighting
% To see the distribution of class labels in the data set, use the countEachLabel function.

tbl = countEachLabel(pxds);
tbl(:,{'Name','PixelCount','ImagePixelCount'})

% The classes in this data set are imbalanced, which is a common issue in automotive data sets containing street scenes. The background class covers more area than the car and truck classes. If not handled correctly, this imbalance can be detrimental to the learning process because the learning is biased in favor of the dominant classes.
% Use these weights to correct the class imbalance. Use the pixel label counts from the tbl.PixelCount property and calculate the median frequency class weights.

imageFreq = tbl.PixelCount ./ tbl.ImagePixelCount;
classWeights = median(imageFreq) ./ imageFreq

%
%% Define Network Architecture
% Create a PointSeg network using the createPointSeg supporting function, which is attached to the example. The code returns the layer graph that you use to train the network.

inputSize = [64 1024 5];

lgraph = createPointSeg(inputSize, classNames, classWeights);

%
% Use the analyzeNetwork function to display an interactive visualization of the network architecture.

analyzeNetwork(lgraph)

%
%% Specify Training Options
% Use the rmsprop optimization algorithm to train the network. Specify the hyperparameters for the algorithm by using the trainingOptions function.

maxEpochs = 30;
initialLearningRate= 5e-4;
miniBatchSize = 8;
l2reg = 2e-4;

options = trainingOptions('rmsprop', ...
    'InitialLearnRate', initialLearningRate, ...
    'L2Regularization', l2reg, ...
    'MaxEpochs', maxEpochs, ...
    'MiniBatchSize', miniBatchSize, ...
    'LearnRateSchedule', 'piecewise', ...
    'LearnRateDropFactor', 0.1, ...
    'LearnRateDropPeriod', 10, ...
    'ValidationData', validationData, ...
    'Plots', 'training-progress', ...
    'VerboseFrequency', 20);

% Note: Reduce miniBatchSize to control memory usage when training.
%
%% Train Network
% Use the trainNetwork function to train a PointSeg network if doTraining is true. Otherwise, load the pretrained network.
% If you train the network, you can use a CPU or a GPU. Using a GPU requires Parallel Computing Toolbox™ and a CUDA® enabled NVIDIA® GPU. For more information, see GPU Support By Release.

if doTraining    
    [net, info] = trainNetwork(trainingData, lgraph, options);
else
    pretrainedNetwork = load('trainedPointSegNet.mat');
    net = pretrainedNetwork.net;
end

%
%% Predict Results on Test Point Cloud
% Use the trained network to predict results on a test point cloud and display the segmentation result.
% First, read a PCD file and convert the point cloud to a five-channel input image. Predict the labels using the trained network. Display the figure with the segmentation as an overlay.

ptCloud = pcread('ousterLidarDrivingData.pcd');
I = helperPointCloudToImage(ptCloud);
predictedResult = semanticseg(I, net);

figure;
helperDisplayLidarOverlayImage(I, predictedResult, classNames);
title('Semantic Segmentation Result');

%
% Use the helperDisplayLidarOverlayPointCloud helper function, defined in the Supporting Functions section of this example, to display the segmentation result over the 3-D point cloud object ptCloud .

figure;
helperDisplayLidarOverlayPointCloud(ptCloud, predictedResult, numClasses);
view([95.71 24.14])
title('Semantic Segmentation Result on Point Cloud');

%
%% Evaluate Network
% Run the semanticseg function on the entire test set to measure the accuracy of the network. Set MiniBatchSize to a value of 8 to reduce memory usage when segmenting images. You can increase or decrease this value depending on the amount of GPU memory you have on your system.

outputLocation = fullfile(tempdir, 'output');
if ~exist(outputLocation,'dir')
    mkdir(outputLocation);
end
pxdsResults = semanticseg(imdsTest, net, ...
                'MiniBatchSize', 8, ...
                'WriteLocation', outputLocation, ...
                'Verbose', false);

%
% The semanticseg function returns the segmentation results on the test data set as a PixelLabelDatastore object. The function writes the actual pixel label data for each test image in the imdsTest object to the disk in the location specified by the 'WriteLocation' argument.
% Use the evaluateSemanticSegmentation function to compute the semantic segmentation metrics from the test set results.

metrics = evaluateSemanticSegmentation(pxdsResults, pxdsTest, 'Verbose', false);

%
% You can measure the amount of overlap per class using the intersection-over-union (IoU) metric.
% The evaluateSemanticSegmentation function returns metrics for the entire data set, for individual classes, and for each test image. To see the metrics at the data set level, use the metrics.DataSetMetrics property.

metrics.DataSetMetrics

% The data set metrics provide a high-level overview of network performance. To see the impact each class has on the overall performance, inspect the metrics for each class using the metrics.ClassMetrics property.

metrics.ClassMetrics

% Although the network overall performance is good, the class metrics show that biased classes (car and truck) are not segmented as well as the classes with abundant data (background). You can improve the network performance by training the network on more labeled data containing the car and truck classes.
%
%% Supporting Functions
%% Function to Augment Data
% The augmentData function randomly flips the 2-D spherical image and associated labels in the horizontal direction.

function out = augmentData(inp)
%augmentData Apply random horizontal flipping.

out = cell(size(inp));

% Randomly flip the five-channel image and pixel labels horizontally.
I = inp{1};
sz = size(I);
tform = randomAffine2d('XReflection',true);
rout = affineOutputView(sz,tform,'BoundsStyle','centerOutput');

out{1} = imwarp(I,tform,'OutputView',rout);
out{2} = imwarp(inp{2},tform,'OutputView',rout);
end

%% Function to Display Lidar Segmentation Map Overlaid on 2-D Spherical Image
% The helperDisplayLidarOverlayImage function overlays the semantic segmentation map over the intensity channel of the 2-D spherical image. The function also resizes the overlaid image for better visualization.

function helperDisplayLidarOverlayImage(lidarImage, labelMap, classNames)
%helperDisplayLidarOverlayImage Overlay labels over the intensity image. 
% 
%  helperDisplayLidarOverlayImage(lidarImage, labelMap, classNames) 
%  displays the overlaid image. lidarImage is a five-channel lidar input. 
%  labelMap contains pixel labels and classNames is an array of label 
%  names.

% Read the intensity channel from the lidar image.
intensityChannel = uint8(lidarImage(:,:,4));

% Load the lidar color map.
cmap = helperLidarColorMap();

% Overlay the labels over the intensity image.
B = labeloverlay(intensityChannel,labelMap,'Colormap',cmap,'Transparency',0.4);

% Resize for better visualization.
B = imresize(B, 'Scale', [3 1], 'method', 'nearest');
imshow(B);

% Display the color bar.
helperPixelLabelColorbar(cmap, classNames); 
end

%% Function To Display Lidar Segmentation Map Overlaid on 3-D Point Cloud
% The helperDisplayLidarOverPointCloud function overlays the segmentation result over a 3-D organized point cloud.

function helperDisplayLidarOverlayPointCloud(ptCloud, labelMap, numClasses)
%helperDisplayLidarOverlayPointCloud Overlay labels over a point cloud object. 
% 
%  helperDisplayLidarOverlayPointCloud(ptCloud, labelMap, numClasses)
%  displays the overlaid pointCloud object. ptCloud is the organized
%  3-D point cloud input. labelMap contains pixel labels and numClasses
%  is the number of predicted classes.

sz = size(labelMap);

% Apply the color red to cars.
carClassCar = zeros(sz(1), sz(2), numClasses, 'uint8');
carClassCar(:,:,1) = 255*ones(sz(1), sz(2), 'uint8');

% Apply the color blue to trucks.
truckClassColor = zeros(sz(1), sz(2), numClasses, 'uint8');
truckClassColor(:,:,3) = 255*ones(sz(1), sz(2), 'uint8');

% Apply the color gray to the background.
backgroundClassColor = 153*ones(sz(1), sz(2), numClasses, 'uint8');

% Extract indices from the labels.
carIndices = labelMap == 'car';
truckIndices = labelMap == 'truck';
backgroundIndices = labelMap == 'background';

% Extract a point cloud for each class.
carPointCloud = select(ptCloud, carIndices, 'OutputSize','full');
truckPointCloud = select(ptCloud, truckIndices, 'OutputSize','full');
backgroundPointCloud = select(ptCloud, backgroundIndices, 'OutputSize','full');

% Apply colors to different classes.
carPointCloud.Color = carClassCar;
truckPointCloud.Color = truckClassColor;
backgroundPointCloud.Color = backgroundClassColor;

% Merge and add all the processed point clouds with class information.
coloredCloud = pcmerge(carPointCloud, truckPointCloud, 0.01);
coloredCloud = pcmerge(coloredCloud, backgroundPointCloud, 0.01);

% Plot the colored point cloud. Set an ROI for better visualization.
ax = pcshow(coloredCloud);
set(ax,'XLim',[-35.0 35.0],'YLim',[-32.0 32.0],'ZLim',[-3 8], ...
    'XColor','none','YColor','none','ZColor','none');
set(get(ax,'parent'), 'units','normalized');
end

%% Function to Define Lidar Colormap
% The helperLidarColorMap function defines the colormap used by the lidar data set.

function cmap = helperLidarColorMap()

cmap = [
   0.00  0.00   0.00  % background
   0.98  0.00   0.00  % car
   0.00  0.00   0.98  % truck
   ];
end

%% Function to Display Pixel Label Colorbar
% The helperPixelLabelColorbar function adds a colorbar to the current axis. The colorbar is formatted to display the class names with the color.

function helperPixelLabelColorbar(cmap, classNames)

colormap(gca, cmap);

% Add a colorbar to the current figure.
c = colorbar('peer', gca);

% Use class names for tick marks.
c.TickLabels = classNames;
numClasses = size(classNames, 1);

% Center tick labels.
c.Ticks = 1/(numClasses * 2):1/numClasses:1;

% Remove tick marks.
c.TickLength = 0;
end

%% References
% [1] Wang, Yuan, Tianyue Shi, Peng Yun, Lei Tai, and Ming Liu. “PointSeg: Real-Time Semantic Segmentation Based on 3D LiDAR Point Cloud.” ArXiv:1807.06288 [Cs], September 25, 2018. http://arxiv.org/abs/1807.06288.
% Copyright 2020 The MathWorks, Inc.