%% Code Generation For Aerial Lidar Semantic Segmentation Using PointNet++ Deep Learning
% This example shows how to generate CUDA® MEX code for a PointNet++ [1] network for lidar semantic segmentation. This example uses a pretrained PointNet++ network that can segment unorganized lidar point clouds belonging to eight classes (buildings, cars, trucks, poles, power lines, fences, ground, and vegetation). For more information on PointNet++ network, see Getting Started with PointNet++.
%
%% Third-Party Prerequisites
%% Required
% CUDA enabled NVIDIA® GPU and compatible driver.
%% Optional
% For non-MEX builds such as static libraries, dynamic libraries, or executables, this example has the following additional requirements.
% NVIDIA toolkit.
% NVIDIA cuDNN library.
% Environment variables for the compilers and libraries. For more information, see Third-party Products and Setting Up the Prerequisite Products.
%
%% Verify GPU Environment
% To verify that the compilers and libraries for running this example are set up correctly, use the coder.checkGpuInstall function.

envCfg = coder.gpuEnvConfig('host');
envCfg.DeepLibTarget = 'cudnn';
envCfg.DeepCodegen = 1;
envCfg.Quiet = 1;
coder.checkGpuInstall(envCfg);

%
%% Load PointNet++ Network
% Use the getPointnetplusNet function, attached as a supporting file to this example, to load the pretrained PointNet++ network. For more information on how to train this network, see Aerial Lidar Semantic Segmentation Using PointNet++ Deep Learning.

net = getPointnetplusNet;

% The pretrained network is a DAG network. To display an interactive visualization of the network architecture, use the analyzeNetwork function.
% The sampling and grouping layer, and the interpolation layer are implemented using the functionLayer function, that does not support code generation. So, replace the function layers in the network with custom layers that support code generation using the helperReplaceFunctionLayers helper function and save the network as a MAT file with the name pointnetplusCodegenNet.mat.

net = helperReplaceFunctionLayers(net);

%
%% pointnetplusPredict Entry-Point Function
% The pointnetplusPredict entry-point function takes a point cloud data matrix as input and performs prediction on it by using the deep learning network saved in the pointnetplusCodegenNet.mat file. The function loads the network object from the pointnetplusCodegenNet.mat file into a persistent variable mynet and reuses the persistent variable in subsequent prediction calls.

type('pointnetplusPredict.m');

%
%% Generate CUDA MEX Code
% To generate CUDA® code for the pointnetplusPredict entry-point function, create a GPU code configuration object for a MEX target and set the target language to C++. Use the coder.DeepLearningConfig function to create a CuDNN deep learning configuration object and assign it to the DeepLearningConfig property of the GPU code configuration object. Run the codegen command with the size of point cloud data in the input layer of the network, which in this case is [8192 1 3].

cfg = coder.gpuConfig('mex');
cfg.TargetLang = 'C++';
cfg.DeepLearningConfig = coder.DeepLearningConfig(TargetLibrary='cudnn');
codegen -config cfg pointnetplusPredict -args {randn(8192,1,3,'single')} -report

% To generate CUDA® code for the TensorRT target, create and use a TensorRT deep learning configuration object instead of the CuDNN configuration object.
%
%% Segment Aerial Point Cloud Using Generated MEX Code
% The network in this example is trained on the DALES data set [2]. Follow the instructions on the DALES website to download the data set to the folder specified by the dataFolder variable. Create folders to store training and test data.

dataFolder = fullfile(tempdir,'DALES');
testDataFolder = fullfile(dataFolder,'dales_las','test');

% Since the network is trained on downsampled point clouds, to perform segmentation on a test point cloud, first downsample the test point cloud, similar to how training data is downsampled. Perform inference on this downsampled test point cloud to compute prediction labels. Interpolate the prediction labels to obtain prediction labels on the dense point cloud.
% Define numNearestNeighbors and radius to find the nearest points in the downsampled point cloud for each point in the dense point cloud and to perform interpolation effectively. Define the parameters gridSize, numPoints, and maxLabel, used to train the network.

numNearestNeighbors = 20;
radius = 0.05;
gridSize = [50,50];
numPoints = 8192;
maxLabel = 1;

% Store the interpolated labels and the target labels in the DensePredictedLabels, and DenseTargetLabels directories, respectively.

densePcTargetLabelsPath = fullfile(testDataFolder,'DenseTargetLabels');
densePcPredLabelsPath =  fullfile(testDataFolder,'DensePredictedLabels');

% Read the full test point cloud.

lasReader = lasFileReader(fullfile(testDataFolder,'5080_54470.las'));
[pc,attr] = readPointCloud(lasReader,'Attributes','Classification');
labelsDenseTarget = attr.Classification;

% Select only labeled data.
pc = select(pc,labelsDenseTarget~=0);
labelsDenseTarget = labelsDenseTarget(labelsDenseTarget~=0);

% Initialize prediction labels. 
labelsDensePred = zeros(size(labelsDenseTarget));
classNames = [
    "ground"
    "vegetation"
    "cars"
    "trucks"
    "powerlines"
    "fences"
    "poles"
    "buildings"
    ];
numClasses = numel(classNames);

% Calculate the number of non-overlapping grids based on the gridSize, XLimits, and YLimits values of the point cloud.

numGridsX = round(diff(pc.XLimits)/gridSize(1));
numGridsY = round(diff(pc.YLimits)/gridSize(2));
[~,edgesX,edgesY,indx,indy] = histcounts2(pc.Location(:,1),pc.Location(:,2), ...
    [numGridsX,numGridsY],'XBinLimits',pc.XLimits,'YBinLimits',pc.YLimits);
ind = sub2ind([numGridsX,numGridsY],indx,indy);

%
% Iterate over all the non-overlapping grids and predict the labels using the pointnetplusPredict_mex function.

for num=1:numGridsX*numGridsY
    idx = ind==num;
    ptCloudDense = select(pc,idx);
    labelsDense = labelsDenseTarget(idx);

    % Use the helperDownsamplePoints function, attached to this example as a
    % supporting file, to extract a downsampled point cloud from the
    % dense point cloud.
    ptCloudSparse = helperDownsamplePoints(ptCloudDense, ...
        labelsDense,numPoints);

    % Make the spatial extent of the dense point cloud and the sparse point
    % cloud the same.
    limits = [ptCloudDense.XLimits;ptCloudDense.YLimits;ptCloudDense.ZLimits];
    ptCloudSparseLocation = ptCloudSparse.Location;
    ptCloudSparseLocation(1:2,:) = limits(:,1:2)';
    ptCloudSparse = pointCloud(ptCloudSparseLocation,'Color',ptCloudSparse.Color, ...
        'Intensity',ptCloudSparse.Intensity, ...
        'Normal',ptCloudSparse.Normal);

    % Use the helperNormalizePointCloud function, attached to this example as
    % a supporting file, to normalize the point cloud between 0 and 1.
    ptCloudSparseNormalized = helperNormalizePointCloud(ptCloudSparse);
    ptCloudDenseNormalized = helperNormalizePointCloud(ptCloudDense);

    % Use the helperConvertPointCloud function, defined at the end of this
    % example, to convert the point cloud to a cell array and to permute the
    % dimensions of the point cloud to make it compatible with the input layer
    % of the network.
    ptCloudSparseForPrediction = helperConvertPointCloud(ptCloudSparseNormalized);

    % Get the output predictions.
    scoresPred = pointnetplusPredict_mex(ptCloudSparseForPrediction{1,1});
    [~,labelsSparsePred] = max(scoresPred,[],3);
    labelsSparsePred = uint8(labelsSparsePred);

    % Use the helperInterpolate function, attached to this example as a
    % supporting file, to calculate labels for the dense point cloud,
    % using the sparse point cloud and labels predicted on the sparse point cloud.
    interpolatedLabels = helperInterpolate(ptCloudDenseNormalized, ...
        ptCloudSparseNormalized,labelsSparsePred,numNearestNeighbors, ...
        radius,maxLabel,numClasses);

    labelsDensePred(idx) = interpolatedLabels;
end

% For better visualization, select a region of interest from the point cloud data. Modify the limits in the roi variable according to the point cloud data.

roi = [edgesX(5) edgesX(8) edgesY(8) edgesY(11) pc.ZLimits];
indices = findPointsInROI(pc,roi); 
figure;
ax = pcshow(select(pc,indices).Location, labelsDensePred(indices));
axis off;
zoom(ax,1.5);
helperLabelColorbar(ax,classNames);
title('Point Cloud Overlaid with Detected Semantic Labels');

%
%% Supporting Functions
% The helperLabelColorbar function adds a colorbar to the current axis. The colorbar is formatted to display the class names with the color.

function helperLabelColorbar(ax,classNames)
% Colormap for the original classes.
cmap = [[0,0,255];
    [0,255,0];
    [255,192,203];
    [255,255,0];
    [255,0,255];
    [255,165,0];
    [139,0,150];
    [255,0,0]];
cmap = cmap./255;
cmap = cmap(1:numel(classNames),:);
colormap(ax,cmap);

% Add colorbar to current figure.
c = colorbar(ax);
c.Color = 'w';

% Center tick labels and use class names for tick marks.
numClasses = size(classNames, 1);
c.Ticks = 1:1:numClasses;
c.TickLabels = classNames;

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

% The helperConvertPointCloud function converts the point cloud to a cell array and permutes the dimensions of the point cloud to make it compatible with the input layer of the network.

function data = helperConvertPointCloud(data)
if ~iscell(data)
    data = {data};
end
numObservations = size(data,1);
for i = 1:numObservations
    tmp = data{i,1}.Location;
    data{i,1} = permute(tmp,[1,3,2]);
end
end

%% References
% [1] Qi, Charles R., Li Yi, Hao Su, and Leonidas J. Guibas. "PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space." ArXiv:1706.02413 [Cs], June 7, 2017. https://arxiv.org/abs/1706.02413.
% [2] Varney, Nina, Vijayan K. Asari, and Quinn Graehling. "DALES: A Large-Scale Aerial LiDAR Data Set for Semantic Segmentation." ArXiv:2004.11985 [Cs, Stat], April 14, 2020. https://arxiv.org/abs/2004.11985.
% Copyright 2021 The MathWorks, Inc.