point_pillars.h

Go to the documentation of this file.

/******************************************************************************
 * Copyright 2020 The Apollo Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *****************************************************************************/

/*
 * Copyright 2018-2019 Autoware Foundation. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once

// headers in STL
#include <algorithm>
#include <cmath>
#include <iomanip>
#include <limits>
#include <map>
#include <memory>
#include <string>
#include <vector>

// headers in TensorRT
#include "NvInfer.h"
#include "NvOnnxParser.h"
#include "torch/script.h"
#include "torch/torch.h"

// headers in local files
#include "modules/perception/lidar/lib/detector/point_pillars_detection/anchor_mask_cuda.h"
#include "modules/perception/lidar/lib/detector/point_pillars_detection/common.h"
#include "modules/perception/lidar/lib/detector/point_pillars_detection/params.h"
#include "modules/perception/lidar/lib/detector/point_pillars_detection/postprocess_cuda.h"
#include "modules/perception/lidar/lib/detector/point_pillars_detection/preprocess_points.h"
#include "modules/perception/lidar/lib/detector/point_pillars_detection/preprocess_points_cuda.h"
#include "modules/perception/lidar/lib/detector/point_pillars_detection/scatter_cuda.h"

namespace apollo {
namespace perception {
namespace lidar {

// Logger for TensorRT info/warning/errors
class Logger : public nvinfer1::ILogger {
 public:
  explicit Logger(Severity severity = Severity::kWARNING)
      : reportable_severity(severity) {}

  void log(Severity severity, const char* msg) override {
    // suppress messages with severity enum value greater than the reportable
    if (severity > reportable_severity) return;

    switch (severity) {
      case Severity::kINTERNAL_ERROR:
        std::cerr << "INTERNAL_ERROR: ";
        break;
      case Severity::kERROR:
        std::cerr << "ERROR: ";
        break;
      case Severity::kWARNING:
        std::cerr << "WARNING: ";
        break;
      case Severity::kINFO:
        std::cerr << "INFO: ";
        break;
      default:
        std::cerr << "UNKNOWN: ";
        break;
    }
    std::cerr << msg << std::endl;
  }

  Severity reportable_severity;
};

class PointPillars {
 private:
  friend class TestClass;
  static const float kPillarXSize;
  static const float kPillarYSize;
  static const float kPillarZSize;
  static const float kMinXRange;
  static const float kMinYRange;
  static const float kMinZRange;
  static const float kMaxXRange;
  static const float kMaxYRange;
  static const float kMaxZRange;
  static const int kNumClass;
  static const int kMaxNumPillars;
  static const int kMaxNumPointsPerPillar;
  static const int kNumPointFeature;
  static const int kGridXSize;
  static const int kGridYSize;
  static const int kGridZSize;
  static const int kRpnInputSize;
  static const int kNumAnchor;
  static const int kNumOutputBoxFeature;
  static const int kRpnBoxOutputSize;
  static const int kRpnClsOutputSize;
  static const int kRpnDirOutputSize;
  static const int kBatchSize;
  static const int kNumIndsForScan;
  static const int kNumThreads;
  // if you change kNumThreads, need to modify NUM_THREADS_MACRO in
  // common.h
  static const int kNumBoxCorners;
  static const std::vector<int> kAnchorStrides;
  static const std::vector<int> kAnchorRanges;
  static const std::vector<int> kNumAnchorSets;
  static const std::vector<std::vector<float>> kAnchorDxSizes;
  static const std::vector<std::vector<float>> kAnchorDySizes;
  static const std::vector<std::vector<float>> kAnchorDzSizes;
  static const std::vector<std::vector<float>> kAnchorZCoors;
  static const std::vector<std::vector<int>> kNumAnchorRo;
  static const std::vector<std::vector<float>> kAnchorRo;

  // initialize in initializer list
  const bool reproduce_result_mode_;
  const float score_threshold_;
  const float nms_overlap_threshold_;
  const std::string pfe_onnx_file_;
  const std::string rpn_onnx_file_;
  const std::string pfe_torch_file_;
  const std::string scattered_torch_file_;
  const std::string backbone_torch_file_;
  const std::string fpn_torch_file_;
  const std::string bbox_head_torch_file_;
  // end initializer list

  int host_pillar_count_[1];

  float* anchors_px_;
  float* anchors_py_;
  float* anchors_pz_;
  float* anchors_dx_;
  float* anchors_dy_;
  float* anchors_dz_;
  float* anchors_ro_;

  float* box_anchors_min_x_;
  float* box_anchors_min_y_;
  float* box_anchors_max_x_;
  float* box_anchors_max_y_;

  int* dev_x_coors_;
  int* dev_y_coors_;
  float* dev_num_points_per_pillar_;
  int* dev_sparse_pillar_map_;
  int* dev_cumsum_along_x_;
  int* dev_cumsum_along_y_;

  float* dev_pillar_point_feature_;
  float* dev_pillar_coors_;

  float* dev_box_anchors_min_x_;
  float* dev_box_anchors_min_y_;
  float* dev_box_anchors_max_x_;
  float* dev_box_anchors_max_y_;
  int* dev_anchor_mask_;

  void* pfe_buffers_[3];
  void* rpn_buffers_[4];

  float* dev_scattered_feature_;

  float* dev_anchors_px_;
  float* dev_anchors_py_;
  float* dev_anchors_pz_;
  float* dev_anchors_dx_;
  float* dev_anchors_dy_;
  float* dev_anchors_dz_;
  float* dev_anchors_ro_;
  float* dev_filtered_box_;
  float* dev_filtered_score_;
  int* dev_filtered_label_;
  int* dev_filtered_dir_;
  float* dev_box_for_nms_;
  int* dev_filter_count_;

  std::unique_ptr<PreprocessPoints> preprocess_points_ptr_;
  std::unique_ptr<PreprocessPointsCuda> preprocess_points_cuda_ptr_;
  std::unique_ptr<AnchorMaskCuda> anchor_mask_cuda_ptr_;
  std::unique_ptr<PostprocessCuda> postprocess_cuda_ptr_;

  Logger g_logger_;
  nvinfer1::ICudaEngine* pfe_engine_;
  nvinfer1::ICudaEngine* rpn_engine_;
  nvinfer1::IExecutionContext* pfe_context_;
  nvinfer1::IExecutionContext* rpn_context_;

  int device_id_ = -1;
  int gpu_id_ = 0;
  torch::DeviceType device_type_;
  torch::jit::script::Module pfe_net_;
  torch::jit::script::Module scattered_net_;
  torch::jit::script::Module backbone_net_;
  torch::jit::script::Module fpn_net_;
  torch::jit::script::Module bbox_head_net_;

  void DeviceMemoryMalloc();

  void InitAnchors();

  void InitTorch();

  void InitTRT();

  void GenerateAnchors(float* anchors_px_, float* anchors_py_,
                       float* anchors_pz_, float* anchors_dx_,
                       float* anchors_dy_, float* anchors_dz_,
                       float* anchors_ro_);

  void OnnxToTRTModel(const std::string& model_file,
                      nvinfer1::ICudaEngine** engine_ptr);

  void Preprocess(const float* in_points_array, const int in_num_points);

  void PreprocessCPU(const float* in_points_array, const int in_num_points);

  void PreprocessGPU(const float* in_points_array, const int in_num_points);

  void ConvertAnchors2BoxAnchors(float* anchors_px_, float* anchors_py_,
                                 float* box_anchors_min_x_,
                                 float* box_anchors_min_y_,
                                 float* box_anchors_max_x_,
                                 float* box_anchors_max_y_);

  void PutAnchorsInDeviceMemory();

 public:
  PointPillars(const bool reproduce_result_mode, const float score_threshold,
               const float nms_overlap_threshold,
               const std::string& pfe_torch_file,
               const std::string& scattered_torch_file,
               const std::string& backbone_torch_file,
               const std::string& fpn_torch_file,
               const std::string& bbox_head_torch_file);
  ~PointPillars();

  void DoInference(const float* in_points_array, const int in_num_points,
                   std::vector<float>* out_detections,
                   std::vector<int>* out_labels);
};

}  // namespace lidar
}  // namespace perception
}  // namespace apollo