%% Code Generation For Lidar Object Detection Using PointPillars Deep Learning
% This example shows how to generate CUDA® MEX for a PointPillars object detector with custom layers. For more information, see Lidar 3-D Object Detection Using PointPillars Deep Learning example from the Lidar Toolbox™.
%
%% Third-Party Prerequisites
%% Required
% CUDA enabled NVIDIA® GPU and compatible driver.
%% Optional
% For non-MEX builds such as static and dynamic libraries or executables, this example has the following additional requirements.
% NVIDIA CUDA 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);

%
%% Pretrained PointPillars Network
% Load the pretrained PointPillars detector trained in the Lidar 3-D Object Detection Using PointPillars Deep Learning example. To train the network yourself, see Lidar 3-D Object Detection Using PointPillars Deep Learning.

pretrainedDetector = load('pretrainedPointPillarsDetector.mat','detector');
detector = pretrainedDetector.detector;

% Extract the pretrained network.

net = detector.Network;

% The scatter layer functionality is implemented using the functionLayer. But, the functionLayer does not support code generation. Use helperReplaceFunctionLayer helper function to replace the functionLayer in the network with the helperScatterLayer custom layer that has code generation support and save the network as pointPillarsCodegenNet.mat file.

[net, matFile] = helperReplaceFunctionLayer(net);

%
%% pointpillarsDetect Entry-Point Function
% The pointpillarsDetect entry-point function takes pillar features and pillar indices as input and passes them to a trained network for prediction through the pointpillarPredict function. The pointpillarsDetect function loads the network object from the MAT file into a persistent variable and reuses the persistent object for subsequent prediction calls. Specifically, the function uses the dlnetwork representation of the network trained in the Lidar 3-D Object Detection Using PointPillars Deep Learning example.

type('pointpillarsDetect.m')

%
%% Evaluate the Entry-Point Function for Object Detection
% Evaluate the entry-point function on a point cloud.
% Define the grid parameters and pillar extraction parameters.
% Read a point cloud and generate pillar features and indices using the createPillars helper function, attached to this example as a supporting file.
% Specify the confidence threshold as 0.7 to keep only detections with confidence scores greater than this value.
% Specify the overlap threshold as 0.1 to remove overlapping detections.
% Use the entry-point function pointpillarsDetect to get the predicted bounding boxes, confidence scores, and class labels.
% Display the point cloud with bounding boxes.
% Define grid parameters.

xMin = detector.PointCloudRange(1,1);   % Minimum value along X-axis.
xMax = detector.PointCloudRange(1,2);   % Maximum value along X-axis.
yMin = detector.PointCloudRange(1,3);   % Minimum value along Y-axis.
yMax = detector.PointCloudRange(1,4);   % Maximum value along Y-axis.
zMin = detector.PointCloudRange(1,5);   % Minimum value along Z-axis.
zMax = detector.PointCloudRange(1,6);   % Maximum value along Z-axis.
xStep = detector.VoxelSize(1,1);   % Resolution along X-axis.
yStep = detector.VoxelSize(1,2);   % Resolution along Y-axis.
dsFactor = 2.0; % DownSampling factor.

% Calculate the dimensions for the pseudo-image.

Xn = round(((xMax - xMin) / xStep));
Yn = round(((yMax - yMin) / yStep));
gridParams = {{xMin,yMin,zMin},{xMax,yMax,zMax},{xStep,yStep,dsFactor},{Xn,Yn}};

% Define the pillar extraction parameters.

P = detector.NumPillars;   % Define number of prominent pillars.
N = detector.NumPointsPerPillar;   % Define number of points per pillar.

% Calculate the number of network outputs.

networkOutputs = numel(net.OutputNames);

% Read a point cloud from the Pandaset data set [2] and convert it to M-by-4 format.

pc = pcread('pandasetDrivingData.pcd');
ptCloud = cat(2,pc.Location,pc.Intensity);

% Create pillar features and pillar indices from the point clouds using the createPillars helper function, attached to this example as a supporting file.

processedPtCloud = createPillars({ptCloud,'',''}, gridParams,P,N);

%  Extract the pillar features and pillar indices.

pillarFeatures = dlarray(single(processedPtCloud{1,1}),'SSCB');
pillarIndices = dlarray(single(processedPtCloud{1,2}),'SSCB');

% Specify the class names.

classNames = cellstr(detector.ClassNames);

% Define the desired thresholds.

confidenceThreshold = 0.7;
overlapThreshold = 0.1;

% Use the detect method on the PointPillars network and display the results.

[bboxes,~,labels] = pointpillarsDetect(matFile, pillarFeatures, pillarIndices,...
    gridParams, networkOutputs, confidenceThreshold, overlapThreshold, classNames);
bboxesCar = bboxes(labels == 'Car',:);
bboxesTruck = bboxes(labels == 'Truck',:);

% Display the detections on the point cloud.

helperDisplay3DBoxesOverlaidPointCloud(pc.Location, bboxesCar, 'green',...
    bboxesTruck, 'magenta', 'Predicted bounding boxes');

%
%% Generate CUDA MEXscatter
% To generate CUDA® code for the pointpillarsDetect 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.

cfg = coder.gpuConfig('mex');
cfg.TargetLang = 'C++';
cfg.DeepLearningConfig = coder.DeepLearningConfig(TargetLibrary='cudnn');

args = {coder.Constant(matFile), pillarFeatures, pillarIndices, gridParams,...
    coder.Constant(networkOutputs), confidenceThreshold, overlapThreshold, classNames};

codegen -config cfg pointpillarsDetect -args args -report

%
%% Run the Generated MEX
% Call the generated CUDA MEX with the same pillar features and indices as before. Display the results.

[bboxes, ~, labels] = pointpillarsDetect_mex(matFile, pillarFeatures, ...
    pillarIndices, gridParams, networkOutputs, confidenceThreshold,...
    overlapThreshold, classNames);
bboxesCar = bboxes(labels == 'Car',:);
bboxesTruck = bboxes(labels == 'Truck',:);

helperDisplay3DBoxesOverlaidPointCloud(pc.Location, bboxesCar, 'green',...
    bboxesTruck, 'magenta', 'Predicted bounding boxes');

%
%% Helper Functions

function helperDisplay3DBoxesOverlaidPointCloud(ptCld, labelsCar, carColor,...
    labelsTruck, truckColor, titleForFigure)
% Display the point cloud with different colored bounding boxes for different
% classes
    figure;
    ax = pcshow(ptCld);
    showShape('cuboid', labelsCar, 'Parent', ax, 'Opacity', 0.1, 'Color',...
        carColor, 'LineWidth', 0.5);
    hold on;
    showShape('cuboid', labelsTruck, 'Parent', ax, 'Opacity', 0.1, 'Color',...
        truckColor, 'LineWidth', 0.5);
    title(titleForFigure);
    zoom(ax,1.5);
end

function [net, matFile] = helperReplaceFunctionLayer(net)
% Replace the scatter functionLayer present in the pretrained Pointpillars
% network with custom layer having code generation support.

    matFile = './pointPillarsCodegenNet.mat';
    lgraph = net.layerGraph;

    id = find(...
        arrayfun( @(x)isa(x,'nnet.cnn.layer.FunctionLayer'), ...
        net.Layers));
    fcnLayer = net.Layers(id);

    customLayer = helperScatterLayer(2,"pillars|scatter_nd",[432 496]);
    lgraph = replaceLayer(lgraph,fcnLayer.Name,customLayer);

    net = dlnetwork(lgraph);
    save(matFile,'net');
end

%% References
% [1] Lang, Alex H., Sourabh Vora, Holger Caesar, Lubing Zhou, Jiong Yang, and Oscar Beijbom. "PointPillars: Fast Encoders for Object Detection From Point Clouds." In 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 12689-12697. Long Beach, CA, USA: IEEE, 2019. https://doi.org/10.1109/CVPR.2019.01298.
% [2] Hesai and Scale. PandaSet. https://scale.com/open-datasets/pandaset.
% Copyright 2020-21 The MathWorks, Inc.