text: import time
  from enum import Enum
  from functools import reduce
  import contextlib
  import numpy as np
  import torch
  from torch import nn
  from torch.nn import functional as F
  
  import torchplus
  from second.pytorch.core import box_torch_ops
  from second.pytorch.core.losses import (WeightedSigmoidClassificationLoss,
                                          WeightedSmoothL1LocalizationLoss,
                                          WeightedSoftmaxClassificationLoss)
  from second.pytorch.models import middle, pointpillars, rpn, voxel_encoder
  from torchplus import metrics
  from second.pytorch.utils import torch_timer
  import mayavi.mlab as mlab
  
  def draw_lidar(pc, color=None, fig=None, bgcolor=(0, 0, 0), pts_scale=1, pts_mode='point', pts_color=None):
      if fig is None: fig = mlab.figure(figure=None, bgcolor=bgcolor, fgcolor=None, engine=None, size=(1600, 1000))
      if color is None: color = pc[:, 2]
      mlab.points3d(pc[:, 0], pc[:, 1], pc[:, 2], color, color=pts_color, mode=pts_mode, colormap='gnuplot',
                    scale_factor=pts_scale, figure=fig)
  
      # draw origin
      mlab.points3d(0, 0, 0, color=(1, 1, 1), mode='sphere', scale_factor=0.2)
  
      # draw axis
      axes = np.array([
          [2., 0., 0., 0.],
          [0., 2., 0., 0.],
          [0., 0., 2., 0.],
      ], dtype=np.float64)
      mlab.plot3d([0, axes[0, 0]], [0, axes[0, 1]], [0, axes[0, 2]], color=(1, 0, 0), tube_radius=None, figure=fig)
      mlab.plot3d([0, axes[1, 0]], [0, axes[1, 1]], [0, axes[1, 2]], color=(0, 1, 0), tube_radius=None, figure=fig)
      mlab.plot3d([0, axes[2, 0]], [0, axes[2, 1]], [0, axes[2, 2]], color=(0, 0, 1), tube_radius=None, figure=fig)
  
      # draw fov (todo: update to real sensor spec.)
      fov = np.array([  # 45 degree
          [20., 20., 0., 0.],
          [20., -20., 0., 0.],
      ], dtype=np.float64)
  
      mlab.plot3d([0, fov[0, 0]], [0, fov[0, 1]], [0, fov[0, 2]], color=(1, 1, 1), tube_radius=None, line_width=1,
                  figure=fig)
      mlab.plot3d([0, fov[1, 0]], [0, fov[1, 1]], [0, fov[1, 2]], color=(1, 1, 1), tube_radius=None, line_width=1,
                  figure=fig)
  
      # draw square region
      TOP_Y_MIN = -20
      TOP_Y_MAX = 20
      TOP_X_MIN = 0
      TOP_X_MAX = 40
      TOP_Z_MIN = -2.0
      TOP_Z_MAX = 0.4
  
      x1 = TOP_X_MIN
      x2 = TOP_X_MAX
      y1 = TOP_Y_MIN
      y2 = TOP_Y_MAX
      mlab.plot3d([x1, x1], [y1, y2], [0, 0], color=(0.5, 0.5, 0.5), tube_radius=0.1, line_width=1, figure=fig)
      mlab.plot3d([x2, x2], [y1, y2], [0, 0], color=(0.5, 0.5, 0.5), tube_radius=0.1, line_width=1, figure=fig)
      mlab.plot3d([x1, x2], [y1, y1], [0, 0], color=(0.5, 0.5, 0.5), tube_radius=0.1, line_width=1, figure=fig)
      mlab.plot3d([x1, x2], [y2, y2], [0, 0], color=(0.5, 0.5, 0.5), tube_radius=0.1, line_width=1, figure=fig)
  
      # mlab.orientation_axes()
      mlab.view(azimuth=180, elevation=70, focalpoint=[12.0909996, -1.04700089, -2.03249991], distance=62.0, figure=fig)
      # mlab.show()
      return fig
  
  def _get_pos_neg_loss(cls_loss, labels):
      # cls_loss: [N, num_anchors, num_class]
      # labels: [N, num_anchors]
      batch_size = cls_loss.shape[0]
      if cls_loss.shape[-1] == 1 or len(cls_loss.shape) == 2:
          cls_pos_loss = (labels > 0).type_as(cls_loss) * cls_loss.view(
              batch_size, -1)
          cls_neg_loss = (labels == 0).type_as(cls_loss) * cls_loss.view(
              batch_size, -1)
          cls_pos_loss = cls_pos_loss.sum() / batch_size
          cls_neg_loss = cls_neg_loss.sum() / batch_size
      else:
          cls_pos_loss = cls_loss[..., 1:].sum() / batch_size
          cls_neg_loss = cls_loss[..., 0].sum() / batch_size
      return cls_pos_loss, cls_neg_loss
  
  REGISTERED_NETWORK_CLASSES = {}
  
  def register_voxelnet(cls, name=None):
      global REGISTERED_NETWORK_CLASSES
      if name is None:
          name = cls.__name__
      assert name not in REGISTERED_NETWORK_CLASSES, f"exist class: {REGISTERED_NETWORK_CLASSES}"
      REGISTERED_NETWORK_CLASSES[name] = cls
      return cls
  
  def get_voxelnet_class(name):
      global REGISTERED_NETWORK_CLASSES
      assert name in REGISTERED_NETWORK_CLASSES, f"available class: {REGISTERED_NETWORK_CLASSES}"
      return REGISTERED_NETWORK_CLASSES[name]
  
  class LossNormType(Enum):
      NormByNumPositives = "norm_by_num_positives"
      NormByNumExamples = "norm_by_num_examples"
      NormByNumPosNeg = "norm_by_num_pos_neg"
      DontNorm = "dont_norm"
  
  @register_voxelnet
  class VoxelNet(nn.Module):
      def __init__(self,
                   output_shape,
                   num_class=2,
                   num_input_features=4,
                   vfe_class_name="VoxelFeatureExtractor",
                   vfe_num_filters=[32, 128],
                   with_distance=False,
                   middle_class_name="SparseMiddleExtractor",
                   middle_num_input_features=-1,
                   middle_num_filters_d1=[64],
                   middle_num_filters_d2=[64, 64],
                   rpn_class_name="RPN",
                   rpn_num_input_features=-1,
                   rpn_layer_nums=[3, 5, 5],
                   rpn_layer_strides=[2, 2, 2],
                   rpn_num_filters=[128, 128, 256],
                   rpn_upsample_strides=[1, 2, 4],
                   rpn_num_upsample_filters=[256, 256, 256],
                   use_norm=True,
                   use_groupnorm=False,
                   num_groups=32,
                   use_direction_classifier=True,
                   use_sigmoid_score=False,
                   encode_background_as_zeros=True,
                   use_rotate_nms=True,
                   multiclass_nms=False,
                   nms_score_thresholds=None,
                   nms_pre_max_sizes=None,
                   nms_post_max_sizes=None,
                   nms_iou_thresholds=None,
                   target_assigner=None,
                   cls_loss_weight=1.0,
                   loc_loss_weight=1.0,
                   pos_cls_weight=1.0,
                   neg_cls_weight=1.0,
                   direction_loss_weight=1.0,
                   loss_norm_type=LossNormType.NormByNumPositives,
                   encode_rad_error_by_sin=False,
                   loc_loss_ftor=None,
                   cls_loss_ftor=None,
                   measure_time=False,
                   voxel_generator=None,
                   post_center_range=None,
                   dir_offset=0.0,
                   sin_error_factor=1.0,
                   nms_class_agnostic=False,
                   num_direction_bins=2,
                   direction_limit_offset=0,
                   name='voxelnet'):
          super().__init__()
          self.name = name
          self._sin_error_factor = sin_error_factor
          self._num_class = num_class
          self._use_rotate_nms = use_rotate_nms
          self._multiclass_nms = multiclass_nms
          self._nms_score_thresholds = nms_score_thresholds
          self._nms_pre_max_sizes = nms_pre_max_sizes
          self._nms_post_max_sizes = nms_post_max_sizes
          self._nms_iou_thresholds = nms_iou_thresholds
          self._use_sigmoid_score = use_sigmoid_score
          self._encode_background_as_zeros = encode_background_as_zeros
          self._use_direction_classifier = use_direction_classifier
          self._num_input_features = num_input_features
          self._box_coder = target_assigner.box_coder
          self.target_assigner = target_assigner
          self.voxel_generator = voxel_generator
          self._pos_cls_weight = pos_cls_weight
          self._neg_cls_weight = neg_cls_weight
          self._encode_rad_error_by_sin = encode_rad_error_by_sin
          self._loss_norm_type = loss_norm_type
          self._dir_loss_ftor = WeightedSoftmaxClassificationLoss()
          self._diff_loc_loss_ftor = WeightedSmoothL1LocalizationLoss()
          self._dir_offset = dir_offset
          self._loc_loss_ftor = loc_loss_ftor
          self._cls_loss_ftor = cls_loss_ftor
          self._direction_loss_weight = direction_loss_weight
          self._cls_loss_weight = cls_loss_weight
          self._loc_loss_weight = loc_loss_weight
          self._post_center_range = post_center_range or []
          self.measure_time = measure_time
          self._nms_class_agnostic = nms_class_agnostic
          self._num_direction_bins = num_direction_bins
          self._dir_limit_offset = direction_limit_offset
          self.voxel_feature_extractor = voxel_encoder.get_vfe_class(vfe_class_name)(
              num_input_features,
              use_norm,
              num_filters=vfe_num_filters,
              with_distance=with_distance,
              voxel_size=self.voxel_generator.voxel_size,
              pc_range=self.voxel_generator.point_cloud_range,
          )
          self.middle_feature_extractor = middle.get_middle_class(middle_class_name)(
              output_shape,
              use_norm,
              num_input_features=middle_num_input_features,
              num_filters_down1=middle_num_filters_d1,
              num_filters_down2=middle_num_filters_d2)
          self.rpn = rpn.get_rpn_class(rpn_class_name)(
              use_norm=True,
              num_class=num_class,
              layer_nums=rpn_layer_nums,
              layer_strides=rpn_layer_strides,
              num_filters=rpn_num_filters,
              upsample_strides=rpn_upsample_strides,
              num_upsample_filters=rpn_num_upsample_filters,
              num_input_features=rpn_num_input_features,
              num_anchor_per_loc=target_assigner.num_anchors_per_location,
              encode_background_as_zeros=encode_background_as_zeros,
              use_direction_classifier=use_direction_classifier,
              use_groupnorm=use_groupnorm,
              num_groups=num_groups,
              box_code_size=target_assigner.box_coder.code_size,
              num_direction_bins=self._num_direction_bins)
          self.rpn_acc = metrics.Accuracy(dim=-1, encode_background_as_zeros=encode_background_as_zeros)
          self.rpn_precision = metrics.Precision(dim=-1)
          self.rpn_recall = metrics.Recall(dim=-1)
          self.rpn_metrics = metrics.PrecisionRecall(
              dim=-1,
              thresholds=[0.1, 0.3, 0.5, 0.7, 0.8, 0.9, 0.95],
              use_sigmoid_score=use_sigmoid_score,
              encode_background_as_zeros=encode_background_as_zeros)
  
          self.rpn_cls_loss = metrics.Scalar()
          self.rpn_loc_loss = metrics.Scalar()
          self.rpn_total_loss = metrics.Scalar()
          self.register_buffer("global_step", torch.LongTensor(1).zero_())
  
          self._time_dict = {}
          self._time_total_dict = {}
          self._time_count_dict = {}
  
      def start_timer(self, *names):
          if not self.measure_time:
              return
          torch.cuda.synchronize()
          for name in names:
              self._time_dict[name] = time.time()
  
      def end_timer(self, name):
          if not self.measure_time:
              return
          torch.cuda.synchronize()
          time_elapsed = time.time() - self._time_dict[name]
          if name not in self._time_count_dict:
              self._time_count_dict[name] = 1
              self._time_total_dict[name] = time_elapsed
          else:
              self._time_count_dict[name] += 1
              self._time_total_dict[name] += time_elapsed
          self._time_dict[name] = 0
  
      def clear_timer(self):
          self._time_count_dict.clear()
          self._time_dict.clear()
          self._time_total_dict.clear()
  
      @contextlib.contextmanager
      def profiler(self):
          old_measure_time = self.measure_time
          self.measure_time = True
          yield
          self.measure_time = old_measure_time
  
      def get_avg_time_dict(self):
          ret = {}
          for name, val in self._time_total_dict.items():
              count = self._time_count_dict[name]
              ret[name] = val / max(1, count)
          return ret
  
      def update_global_step(self):
          self.global_step += 1
  
      def get_global_step(self):
          return int(self.global_step.cpu().numpy()[0])
  
      def clear_global_step(self):
          self.global_step.zero_()
  
      def loss(self, example, preds_dict):
          box_preds = preds_dict["box_preds"]
          cls_preds = preds_dict["cls_preds"]
          batch_size_dev = cls_preds.shape[0]
          self.start_timer("loss forward")
          labels = example['labels']
          reg_targets = example['reg_targets']
          importance = example['importance']
          self.start_timer("prepare weight forward")
          cls_weights, reg_weights, cared = prepare_loss_weights(
              labels,
              pos_cls_weight=self._pos_cls_weight,
              neg_cls_weight=self._neg_cls_weight,
              loss_norm_type=self._loss_norm_type,
              dtype=box_preds.dtype)
  
          cls_targets = labels * cared.type_as(labels)
          cls_targets = cls_targets.unsqueeze(-1)
          self.end_timer("prepare weight forward")
          self.start_timer("create_loss forward")
          loc_loss, cls_loss = create_loss(
              self._loc_loss_ftor,
              self._cls_loss_ftor,
              box_preds=box_preds,
              cls_preds=cls_preds,
              cls_targets=cls_targets,
              cls_weights=cls_weights * importance,
              reg_targets=reg_targets,
              reg_weights=reg_weights * importance,
              num_class=self._num_class,
              encode_rad_error_by_sin=self._encode_rad_error_by_sin,
              encode_background_as_zeros=self._encode_background_as_zeros,
              box_code_size=self._box_coder.code_size,
              sin_error_factor=self._sin_error_factor,
              num_direction_bins=self._num_direction_bins,
          )
          loc_loss_reduced = loc_loss.sum() / batch_size_dev
          loc_loss_reduced *= self._loc_loss_weight
          cls_pos_loss, cls_neg_loss = _get_pos_neg_loss(cls_loss, labels)
          cls_pos_loss /= self._pos_cls_weight
          cls_neg_loss /= self._neg_cls_weight
          cls_loss_reduced = cls_loss.sum() / batch_size_dev
          cls_loss_reduced *= self._cls_loss_weight
          loss = loc_loss_reduced + cls_loss_reduced
          self.end_timer("create_loss forward")
          if self._use_direction_classifier:
              dir_targets = get_direction_target(
                  example['anchors'],
                  reg_targets,
                  dir_offset=self._dir_offset,
                  num_bins=self._num_direction_bins)
              dir_logits = preds_dict["dir_cls_preds"].view(
                  batch_size_dev, -1, self._num_direction_bins)
              weights = (labels > 0).type_as(dir_logits) * importance
              weights /= torch.clamp(weights.sum(-1, keepdim=True), min=1.0)
              dir_loss = self._dir_loss_ftor(
                  dir_logits, dir_targets, weights=weights)
              dir_loss = dir_loss.sum() / batch_size_dev
              loss += dir_loss * self._direction_loss_weight
          self.end_timer("loss forward")
          res = {
              "loss": loss,
              "cls_loss": cls_loss,
              "loc_loss": loc_loss,
              "cls_pos_loss": cls_pos_loss,
              "cls_neg_loss": cls_neg_loss,
              "cls_preds": cls_preds,
              "cls_loss_reduced": cls_loss_reduced,
              "loc_loss_reduced": loc_loss_reduced,
              "cared": cared,
          }
          if self._use_direction_classifier:
              res["dir_loss_reduced"] = dir_loss
          return res
  
      def network_forward(self, voxels, num_points, coors, batch_size):
          """this function is used for subclass.
          you can add custom network architecture by subclass VoxelNet class
          and override this function.
          Returns:
              preds_dict: {
                  box_preds: ...
                  cls_preds: ...
                  dir_cls_preds: ...
              }
          """
          self.start_timer("voxel_feature_extractor")
          voxel_features = self.voxel_feature_extractor(voxels, num_points,
                                                        coors)
       ########################################################################33
          fig = mlab.figure(figure=None, bgcolor=(0, 0, 0), fgcolor=None, engine=None, size=(1000, 500))
          pcd_data = voxel_features.detach().cpu().numpy()
          draw_lidar(pcd_data, fig=fig)
          mlab.show()
  ####################################################################################
          self.end_timer("voxel_feature_extractor")
  
          self.start_timer("middle forward")
          spatial_features = self.middle_feature_extractor(
              voxel_features, coors, batch_size)
          self.end_timer("middle forward")
          self.start_timer("rpn forward")
          preds_dict = self.rpn(spatial_features)
          self.end_timer("rpn forward")
          return preds_dict
  
      def forward(self, example):
          """module's forward should always accept dict and return loss.
          """
          voxels = example["voxels"]
          num_points = example["num_points"]
          coors = example["coordinates"]
          if len(num_points.shape) == 2:  # multi-gpu
              num_voxel_per_batch = example["num_voxels"].cpu().numpy().reshape(
                  -1)
              voxel_list = []
              num_points_list = []
              coors_list = []
              for i, num_voxel in enumerate(num_voxel_per_batch):
                  voxel_list.append(voxels[i, :num_voxel])
                  num_points_list.append(num_points[i, :num_voxel])
                  coors_list.append(coors[i, :num_voxel])
              voxels = torch.cat(voxel_list, dim=0)
              num_points = torch.cat(num_points_list, dim=0)
              coors = torch.cat(coors_list, dim=0)
          batch_anchors = example["anchors"]
          batch_size_dev = batch_anchors.shape[0]
          # features: [num_voxels, max_num_points_per_voxel, 7]
          # num_points: [num_voxels]
          # coors: [num_voxels, 4]
          preds_dict = self.network_forward(voxels, num_points, coors, batch_size_dev)
          # need to check size.
          box_preds = preds_dict["box_preds"].view(batch_size_dev, -1, self._box_coder.code_size)
          err_msg = f"num_anchors={batch_anchors.shape[1]}, but num_output={box_preds.shape[1]}. please check size"
          assert batch_anchors.shape[1] == box_preds.shape[1], err_msg
          if self.training:
              return self.loss(example, preds_dict)
          else:
              self.start_timer("predict")
              with torch.no_grad():
                  res = self.predict(example, preds_dict)
              self.end_timer("predict")
              return res
  
      def predict(self, example, preds_dict):
          """start with v1.6.0, this function don't contain any kitti-specific code.
          Returns:
              predict: list of pred_dict.
              pred_dict: {
                  box3d_lidar: [N, 7] 3d box.
                  scores: [N]
                  label_preds: [N]
                  metadata: meta-data which contains dataset-specific information.
                      for kitti, it contains image idx (label idx),
                      for nuscenes, sample_token is saved in it.
              }
          """
          batch_size = example['anchors'].shape[0]
          if "metadata" not in example or len(example["metadata"]) == 0:
              meta_list = [None] * batch_size
          else:
              meta_list = example["metadata"]
          batch_anchors = example["anchors"].view(batch_size, -1,
                                                  example["anchors"].shape[-1])
          if "anchors_mask" not in example:
              batch_anchors_mask = [None] * batch_size
          else:
              batch_anchors_mask = example["anchors_mask"].view(batch_size, -1)
  
          t = time.time()
          batch_box_preds = preds_dict["box_preds"]
          batch_cls_preds = preds_dict["cls_preds"]
          batch_box_preds = batch_box_preds.view(batch_size, -1,
                                                 self._box_coder.code_size)
          num_class_with_bg = self._num_class
          if not self._encode_background_as_zeros:
              num_class_with_bg = self._num_class + 1
  
          batch_cls_preds = batch_cls_preds.view(batch_size, -1,
                                                 num_class_with_bg)
          batch_box_preds = self._box_coder.decode_torch(batch_box_preds,
                                                         batch_anchors)
          if self._use_direction_classifier:
              batch_dir_preds = preds_dict["dir_cls_preds"]
              batch_dir_preds = batch_dir_preds.view(batch_size, -1,
                                                     self._num_direction_bins)
          else:
              batch_dir_preds = [None] * batch_size
  
          predictions_dicts = []
          post_center_range = None
          if len(self._post_center_range) > 0:
              post_center_range = torch.tensor(
                  self._post_center_range,
                  dtype=batch_box_preds.dtype,
                  device=batch_box_preds.device).float()
          for box_preds, cls_preds, dir_preds, a_mask, meta in zip(
                  batch_box_preds, batch_cls_preds, batch_dir_preds,
                  batch_anchors_mask, meta_list):
              if a_mask is not None:
                  box_preds = box_preds[a_mask]
                  cls_preds = cls_preds[a_mask]
              box_preds = box_preds.float()
              cls_preds = cls_preds.float()
              if self._use_direction_classifier:
                  if a_mask is not None:
                      dir_preds = dir_preds[a_mask]
                  dir_labels = torch.max(dir_preds, dim=-1)[1]
              if self._encode_background_as_zeros:
                  # this don't support softmax
                  assert self._use_sigmoid_score is True
                  total_scores = torch.sigmoid(cls_preds)
              else:
                  # encode background as first element in one-hot vector
                  if self._use_sigmoid_score:
                      total_scores = torch.sigmoid(cls_preds)[..., 1:]
                  else:
                      total_scores = F.softmax(cls_preds, dim=-1)[..., 1:]
              # Apply NMS in birdeye view
              if self._use_rotate_nms:
                  nms_func = box_torch_ops.rotate_nms
              else:
                  nms_func = box_torch_ops.nms
              feature_map_size_prod = batch_box_preds.shape[
                  1] // self.target_assigner.num_anchors_per_location
              if self._multiclass_nms:
                  assert self._encode_background_as_zeros is True
                  boxes_for_nms = box_preds[:, [0, 1, 3, 4, 6]]
                  if not self._use_rotate_nms:
                      box_preds_corners = box_torch_ops.center_to_corner_box2d(
                          boxes_for_nms[:, :2], boxes_for_nms[:, 2:4],
                          boxes_for_nms[:, 4])
                      boxes_for_nms = box_torch_ops.corner_to_standup_nd(
                          box_preds_corners)
  
                  selected_boxes, selected_labels, selected_scores = [], [], []
                  selected_dir_labels = []
  
                  scores = total_scores
                  boxes = boxes_for_nms
                  selected_per_class = []
                  score_threshs = self._nms_score_thresholds
                  pre_max_sizes = self._nms_pre_max_sizes
                  post_max_sizes = self._nms_post_max_sizes
                  iou_thresholds = self._nms_iou_thresholds
                  for class_idx, score_thresh, pre_ms, post_ms, iou_th in zip(
                          range(self._num_class),
                          score_threshs,
                          pre_max_sizes, post_max_sizes, iou_thresholds):
                      if self._nms_class_agnostic:
                          class_scores = total_scores.view(
                              feature_map_size_prod, -1,
                              self._num_class)[..., class_idx]
                          class_scores = class_scores.contiguous().view(-1)
                          class_boxes_nms = boxes.view(-1,
                                                       boxes_for_nms.shape[-1])
                          class_boxes = box_preds
                          class_dir_labels = dir_labels
                      else:
                          anchors_range = self.target_assigner.anchors_range(class_idx)
                          class_scores = total_scores.view(
                              -1,
                              self._num_class)[anchors_range[0]:anchors_range[1], class_idx]
                          class_boxes_nms = boxes.view(-1,
                              boxes_for_nms.shape[-1])[anchors_range[0]:anchors_range[1], :]
                          class_scores = class_scores.contiguous().view(-1)
                          class_boxes_nms = class_boxes_nms.contiguous().view(
                              -1, boxes_for_nms.shape[-1])
                          class_boxes = box_preds.view(-1,
                              box_preds.shape[-1])[anchors_range[0]:anchors_range[1], :]
                          class_boxes = class_boxes.contiguous().view(
                              -1, box_preds.shape[-1])
                          if self._use_direction_classifier:
                              class_dir_labels = dir_labels.view(-1)[anchors_range[0]:anchors_range[1]]
                              class_dir_labels = class_dir_labels.contiguous(
                              ).view(-1)
                      if score_thresh > 0.0:
                          class_scores_keep = class_scores >= score_thresh
                          if class_scores_keep.shape[0] == 0:
                              selected_per_class.append(None)
                              continue
                          class_scores = class_scores[class_scores_keep]
                      if class_scores.shape[0] != 0:
                          if score_thresh > 0.0:
                              class_boxes_nms = class_boxes_nms[
                                  class_scores_keep]
                              class_boxes = class_boxes[class_scores_keep]
                              class_dir_labels = class_dir_labels[
                                  class_scores_keep]
                          keep = nms_func(class_boxes_nms, class_scores, pre_ms,
                                          post_ms, iou_th)
                          if keep.shape[0] != 0:
                              selected_per_class.append(keep)
                          else:
                              selected_per_class.append(None)
                      else:
                          selected_per_class.append(None)
                      selected = selected_per_class[-1]
  
                      if selected is not None:
                          selected_boxes.append(class_boxes[selected])
                          selected_labels.append(
                              torch.full([class_boxes[selected].shape[0]],
                                         class_idx,
                                         dtype=torch.int64,
                                         device=box_preds.device))
                          if self._use_direction_classifier:
                              selected_dir_labels.append(
                                  class_dir_labels[selected])
                          selected_scores.append(class_scores[selected])
                  selected_boxes = torch.cat(selected_boxes, dim=0)
                  selected_labels = torch.cat(selected_labels, dim=0)
                  selected_scores = torch.cat(selected_scores, dim=0)
                  if self._use_direction_classifier:
                      selected_dir_labels = torch.cat(selected_dir_labels, dim=0)
              else:
                  # get highest score per prediction, than apply nms
                  # to remove overlapped box.
                  if num_class_with_bg == 1:
                      top_scores = total_scores.squeeze(-1)
                      top_labels = torch.zeros(
                          total_scores.shape[0],
                          device=total_scores.device,
                          dtype=torch.long)
                  else:
                      top_scores, top_labels = torch.max(
                          total_scores, dim=-1)
                  if self._nms_score_thresholds[0] > 0.0:
                      top_scores_keep = top_scores >= self._nms_score_thresholds[0]
                      top_scores = top_scores.masked_select(top_scores_keep)
  
                  if top_scores.shape[0] != 0:
                      if self._nms_score_thresholds[0] > 0.0:
                          box_preds = box_preds[top_scores_keep]
                          if self._use_direction_classifier:
                              dir_labels = dir_labels[top_scores_keep]
                          top_labels = top_labels[top_scores_keep]
                      boxes_for_nms = box_preds[:, [0, 1, 3, 4, 6]]
                      if not self._use_rotate_nms:
                          box_preds_corners = box_torch_ops.center_to_corner_box2d(
                              boxes_for_nms[:, :2], boxes_for_nms[:, 2:4],
                              boxes_for_nms[:, 4])
                          boxes_for_nms = box_torch_ops.corner_to_standup_nd(
                              box_preds_corners)
                      # the nms in 3d detection just remove overlap boxes.
                      selected = nms_func(
                          boxes_for_nms,
                          top_scores,
                          pre_max_size=self._nms_pre_max_sizes[0],
                          post_max_size=self._nms_post_max_sizes[0],
                          iou_threshold=self._nms_iou_thresholds[0],
                      )
                  else:
                      selected = []
                  # if selected is not None:
                  selected_boxes = box_preds[selected]
                  if self._use_direction_classifier:
                      selected_dir_labels = dir_labels[selected]
                  selected_labels = top_labels[selected]
                  selected_scores = top_scores[selected]
              # finally generate predictions.
              if selected_boxes.shape[0] != 0:
                  box_preds = selected_boxes
                  scores = selected_scores
                  label_preds = selected_labels
                  if self._use_direction_classifier:
                      dir_labels = selected_dir_labels
                      period = (2 * np.pi / self._num_direction_bins)
                      dir_rot = box_torch_ops.limit_period(
                          box_preds[..., 6] - self._dir_offset,
                          self._dir_limit_offset, period)
                      box_preds[
                          ...,
                          6] = dir_rot + self._dir_offset + period * dir_labels.to(
                              box_preds.dtype)
                  final_box_preds = box_preds
                  final_scores = scores
                  final_labels = label_preds
                  if post_center_range is not None:
                      mask = (final_box_preds[:, :3] >=
                              post_center_range[:3]).all(1)
                      mask &= (final_box_preds[:, :3] <=
                               post_center_range[3:]).all(1)
                      predictions_dict = {
                          "box3d_lidar": final_box_preds[mask],
                          "scores": final_scores[mask],
                          "label_preds": label_preds[mask],
                          "metadata": meta,
                      }
                  else:
                      predictions_dict = {
                          "box3d_lidar": final_box_preds,
                          "scores": final_scores,
                          "label_preds": label_preds,
                          "metadata": meta,
                      }
              else:
                  dtype = batch_box_preds.dtype
                  device = batch_box_preds.device
                  predictions_dict = {
                      "box3d_lidar":
                      torch.zeros([0, box_preds.shape[-1]],
                                  dtype=dtype,
                                  device=device),
                      "scores":
                      torch.zeros([0], dtype=dtype, device=device),
                      "label_preds":
                      torch.zeros([0], dtype=top_labels.dtype, device=device),
                      "metadata":
                      meta,
                  }
              predictions_dicts.append(predictions_dict)
          return predictions_dicts
  
      def metrics_to_float(self):
          self.rpn_acc.float()
          self.rpn_metrics.float()
          self.rpn_cls_loss.float()
          self.rpn_loc_loss.float()
          self.rpn_total_loss.float()
  
      def update_metrics(self, cls_loss, loc_loss, cls_preds, labels, sampled):
          batch_size = cls_preds.shape[0]
          num_class = self._num_class
          if not self._encode_background_as_zeros:
              num_class += 1
          cls_preds = cls_preds.view(batch_size, -1, num_class)
          rpn_acc = self.rpn_acc(labels, cls_preds, sampled).numpy()[0]
          prec, recall = self.rpn_metrics(labels, cls_preds, sampled)
          precall = self.rpn_metrics(labels, cls_preds, sampled)
          prec = prec.numpy()
          recall = recall.numpy()
          rpn_cls_loss = self.rpn_cls_loss(cls_loss).numpy()[0]
          rpn_loc_loss = self.rpn_loc_loss(loc_loss).numpy()[0]
          ret = {
              "loss": {
                  "cls_loss": float(rpn_cls_loss),
                  "cls_loss_rt": float(cls_loss.data.cpu().numpy()),
                  'loc_loss': float(rpn_loc_loss),
                  "loc_loss_rt": float(loc_loss.data.cpu().numpy()),
              },
              "rpn_acc": float(rpn_acc),
              "pr": {},
          }
          for i, thresh in enumerate(self.rpn_metrics.thresholds):
              ret["pr"][f"prec@{int(thresh*100)}"] = float(prec[i])
              ret["pr"][f"rec@{int(thresh*100)}"] = float(recall[i])
          return ret
  
      def clear_metrics(self):
          self.rpn_acc.clear()
          self.rpn_metrics.clear()
          self.rpn_cls_loss.clear()
          self.rpn_loc_loss.clear()
          self.rpn_total_loss.clear()
  
      @staticmethod
      def convert_norm_to_float(net):
          '''
          BatchNorm layers to have parameters in single precision.
          Find all layers and convert them back to float. This can't
          be done with built in .apply as that function will apply
          fn to all modules, parameters, and buffers. Thus we wouldn't
          be able to guard the float conversion based on the module type.
          '''
          if isinstance(net, torch.nn.modules.batchnorm._BatchNorm):
              net.float()
          for child in net.children():
              VoxelNet.convert_norm_to_float(child)
          return net
  
  
  def add_sin_difference(boxes1, boxes2, boxes1_rot, boxes2_rot, factor=1.0):
      if factor != 1.0:
          boxes1_rot = factor * boxes1_rot
          boxes2_rot = factor * boxes2_rot
      rad_pred_encoding = torch.sin(boxes1_rot) * torch.cos(boxes2_rot)
      rad_tg_encoding = torch.cos(boxes1_rot) * torch.sin(boxes2_rot)
      boxes1 = torch.cat([boxes1[..., :6], rad_pred_encoding, boxes1[..., 7:]],
                         dim=-1)
      boxes2 = torch.cat([boxes2[..., :6], rad_tg_encoding, boxes2[..., 7:]],
                         dim=-1)
      return boxes1, boxes2
  
  
  def create_loss(loc_loss_ftor,
                  cls_loss_ftor,
                  box_preds,
                  cls_preds,
                  cls_targets,
                  cls_weights,
                  reg_targets,
                  reg_weights,
                  num_class,
                  encode_background_as_zeros=True,
                  encode_rad_error_by_sin=True,
                  sin_error_factor=1.0,
                  box_code_size=7,
                  num_direction_bins=2):
      batch_size = int(box_preds.shape[0])
      box_preds = box_preds.view(batch_size, -1, box_code_size)
      if encode_background_as_zeros:
          cls_preds = cls_preds.view(batch_size, -1, num_class)
      else:
          cls_preds = cls_preds.view(batch_size, -1, num_class + 1)
      cls_targets = cls_targets.squeeze(-1)
      one_hot_targets = torchplus.nn.one_hot(
          cls_targets, depth=num_class + 1, dtype=box_preds.dtype)
      if encode_background_as_zeros:
          one_hot_targets = one_hot_targets[..., 1:]
      if encode_rad_error_by_sin:
          # sin(a - b) = sinacosb-cosasinb
          # reg_tg_rot = box_torch_ops.limit_period(
          #     reg_targets[..., 6:7], 0.5, 2 * np.pi / num_direction_bins)
          box_preds, reg_targets = add_sin_difference(box_preds, reg_targets, box_preds[..., 6:7], reg_targets[..., 6:7],
                                                      sin_error_factor)
  
      loc_losses = loc_loss_ftor(
          box_preds, reg_targets, weights=reg_weights)  # [N, M]
      cls_losses = cls_loss_ftor(
          cls_preds, one_hot_targets, weights=cls_weights)  # [N, M]
      return loc_losses, cls_losses
  
  
  def prepare_loss_weights(labels,
                           pos_cls_weight=1.0,
                           neg_cls_weight=1.0,
                           loss_norm_type=LossNormType.NormByNumPositives,
                           dtype=torch.float32):
      """get cls_weights and reg_weights from labels.
      """
      cared = labels >= 0
      # cared: [N, num_anchors]
      positives = labels > 0
      negatives = labels == 0
      negative_cls_weights = negatives.type(dtype) * neg_cls_weight
      cls_weights = negative_cls_weights + pos_cls_weight * positives.type(dtype)
      reg_weights = positives.type(dtype)
      if loss_norm_type == LossNormType.NormByNumExamples:
          num_examples = cared.type(dtype).sum(1, keepdim=True)
          num_examples = torch.clamp(num_examples, min=1.0)
          cls_weights /= num_examples
          bbox_normalizer = positives.sum(1, keepdim=True).type(dtype)
          reg_weights /= torch.clamp(bbox_normalizer, min=1.0)
      elif loss_norm_type == LossNormType.NormByNumPositives:  # for focal loss
          pos_normalizer = positives.sum(1, keepdim=True).type(dtype)
          reg_weights /= torch.clamp(pos_normalizer, min=1.0)
          cls_weights /= torch.clamp(pos_normalizer, min=1.0)
      elif loss_norm_type == LossNormType.NormByNumPosNeg:
          pos_neg = torch.stack([positives, negatives], dim=-1).type(dtype)
          normalizer = pos_neg.sum(1, keepdim=True)  # [N, 1, 2]
          cls_normalizer = (pos_neg * normalizer).sum(-1)  # [N, M]
          cls_normalizer = torch.clamp(cls_normalizer, min=1.0)
          # cls_normalizer will be pos_or_neg_weight/num_pos_or_neg
          normalizer = torch.clamp(normalizer, min=1.0)
          reg_weights /= normalizer[:, 0:1, 0]
          cls_weights /= cls_normalizer
      elif loss_norm_type == LossNormType.DontNorm:  # support ghm loss
          pos_normalizer = positives.sum(1, keepdim=True).type(dtype)
          reg_weights /= torch.clamp(pos_normalizer, min=1.0)
      else:
          raise ValueError(
              f"unknown loss norm type. available: {list(LossNormType)}")
      return cls_weights, reg_weights, cared
  
  
  def assign_weight_to_each_class(labels,
                                  weight_per_class,
                                  norm_by_num=True,
                                  dtype=torch.float32):
      weights = torch.zeros(labels.shape, dtype=dtype, device=labels.device)
      for label, weight in weight_per_class:
          positives = (labels == label).type(dtype)
          weight_class = weight * positives
          if norm_by_num:
              normalizer = positives.sum()
              normalizer = torch.clamp(normalizer, min=1.0)
              weight_class /= normalizer
          weights += weight_class
      return weights
  
  
  def get_direction_target(anchors,
                           reg_targets,
                           one_hot=True,
                           dir_offset=0,
                           num_bins=2):
      batch_size = reg_targets.shape[0]
      anchors = anchors.view(batch_size, -1, anchors.shape[-1])
      rot_gt = reg_targets[..., 6] + anchors[..., 6]
      offset_rot = box_torch_ops.limit_period(rot_gt - dir_offset, 0, 2 * np.pi)
      dir_cls_targets = torch.floor(offset_rot / (2 * np.pi / num_bins)).long()
      dir_cls_targets = torch.clamp(dir_cls_targets, min=0, max=num_bins - 1)
      if one_hot:
          dir_cls_targets = torchplus.nn.one_hot(
              dir_cls_targets, num_bins, dtype=anchors.dtype)
      return dir_cls_targets
  