text: from operator import gt
  import torch
  import copy
  import mmcv
  import numpy as np
  import os
  import tempfile
  from mmcv.utils import print_log
  from os import path as osp
  
  from mmdet.datasets import DATASETS
  from .kitti_dataset import KittiDataset
  
  from ..core.bbox import (Box3DMode, CameraInstance3DBoxes, Coord3DMode,
                           LiDARInstance3DBoxes, points_cam2img)
  
  @DATASETS.register_module()
  class CustomKittiDataset(KittiDataset):
  
      CLASSES = ('car', 'pedestrian', 'cyclist')
  
      def __init__(self,
                      data_root,
                      ann_file,
                      split,
                      pts_prefix='velodyne',
                      pipeline=None,
                      classes=None,
                      modality=None,
                      box_type_3d='LiDAR',
                      filter_empty_gt=True,
                      test_mode=False,
                      pcd_limit_range=[0, -40, -3, 70.4, 40, 0.0],
                      custom_range=[0, 50000, 200000, 100000],
                      custom_attr='distance'):
          super().__init__(
              data_root=data_root,
              ann_file=ann_file,
              split=split,
              pipeline=pipeline,
              classes=classes,
              modality=modality,
              box_type_3d=box_type_3d,
              filter_empty_gt=filter_empty_gt,
              test_mode=test_mode,
              pcd_limit_range=pcd_limit_range)
  
          self.split = split
          self.root_split = os.path.join(self.data_root, split)
          assert self.modality is not None
          self.pcd_limit_range = pcd_limit_range
          self.pts_prefix = pts_prefix
          self.custom_range = np.array(custom_range)
          self.custom_attr = custom_attr
      
      def format_results(self,
                         outputs,
                         pklfile_prefix=None,
                         submission_prefix=None):
          """Format the results to pkl file.
  
          Args:
              outputs (list[dict]): Testing results of the dataset.
              pklfile_prefix (str | None): The prefix of pkl files. It includes
                  the file path and the prefix of filename, e.g., "a/b/prefix".
                  If not specified, a temp file will be created. Default: None.
              submission_prefix (str | None): The prefix of submitted files. It
                  includes the file path and the prefix of filename, e.g.,
                  "a/b/prefix". If not specified, a temp file will be created.
                  Default: None.
  
          Returns:
              tuple: (result_files, tmp_dir), result_files is a dict containing \
                  the json filepaths, tmp_dir is the temporal directory created \
                  for saving json files when jsonfile_prefix is not specified.
          """
          if pklfile_prefix is None:
              tmp_dir = tempfile.TemporaryDirectory()
              pklfile_prefix = osp.join(tmp_dir.name, 'results')
          else:
              tmp_dir = None
  
          if not isinstance(outputs[0], dict):
              result_files = self.bbox2result_kitti2d(outputs, self.CLASSES,
                                                      pklfile_prefix,
                                                      submission_prefix)
          elif 'pts_bbox' in outputs[0] or 'img_bbox' in outputs[0]:
              result_files = dict()
              for name in outputs[0]:
                  results_ = [out[name] for out in outputs]
                  pklfile_prefix_ = pklfile_prefix + name
                  if submission_prefix is not None:
                      submission_prefix_ = submission_prefix + name
                  else:
                      submission_prefix_ = None
                  if 'img' in name:
                      result_files = self.bbox2result_kitti2d(
                          results_, self.CLASSES, pklfile_prefix_,
                          submission_prefix_)
                  else:
                      result_files_ = self.bbox2result_kitti(
                          results_, self.CLASSES, pklfile_prefix_,
                          submission_prefix_)
                  result_files[name] = result_files_
          else:
              result_files = self.bbox2result_kitti(outputs, self.CLASSES,
                                                    pklfile_prefix,
                                                    submission_prefix)
          return result_files, tmp_dir
      
      def convert_valid_bboxes(self, box_dict, info):
          """Convert the predicted boxes into valid ones.
  
          Args:
              box_dict (dict): Box dictionaries to be converted.
  
                  - boxes_3d (:obj:`LiDARInstance3DBoxes`): 3D bounding boxes.
                  - scores_3d (torch.Tensor): Scores of boxes.
                  - labels_3d (torch.Tensor): Class labels of boxes.
              info (dict): Data info.
  
          Returns:
              dict: Valid predicted boxes.
  
                  - bbox (np.ndarray): 2D bounding boxes.
                  - box3d_camera (np.ndarray): 3D bounding boxes in \
                      camera coordinate.
                  - box3d_lidar (np.ndarray): 3D bounding boxes in \
                      LiDAR coordinate.
                  - scores (np.ndarray): Scores of boxes.
                  - label_preds (np.ndarray): Class label predictions.
                  - sample_idx (int): Sample index.
          """
          # TODO: refactor this function
          box_preds = box_dict['boxes_3d']
          scores = box_dict['scores_3d']
          labels = box_dict['labels_3d']
          customs_pred = box_dict['attrs_3d']
          sample_idx = info['image']['image_idx']
          # TODO: remove the hack of yaw
          box_preds.tensor[:, -1] = box_preds.tensor[:, -1] - np.pi
          box_preds.limit_yaw(offset=0.5, period=np.pi * 2)
  
          if len(box_preds) == 0:
              return dict(
                  bbox=np.zeros([0, 4]),
                  box3d_camera=np.zeros([0, 7]),
                  box3d_lidar=np.zeros([0, 7]),
                  scores=np.zeros([0]),
                  customs_pred = np.zeros([0]),
                  label_preds=np.zeros([0, 4]),
                  sample_idx=sample_idx)
  
          rect = info['calib']['R0_rect'].astype(np.float32)
          Trv2c = info['calib']['Tr_velo_to_cam'].astype(np.float32)
          P2 = info['calib']['P2'].astype(np.float32)
          img_shape = info['image']['image_shape']
          P2 = box_preds.tensor.new_tensor(P2)
  
          box_preds_camera = box_preds.convert_to(Box3DMode.CAM, rect @ Trv2c)
  
          box_corners = box_preds_camera.corners
          box_corners_in_image = points_cam2img(box_corners, P2)
          # box_corners_in_image: [N, 8, 2]
          minxy = torch.min(box_corners_in_image, dim=1)[0]
          maxxy = torch.max(box_corners_in_image, dim=1)[0]
          box_2d_preds = torch.cat([minxy, maxxy], dim=1)
          # Post-processing
          # check box_preds_camera
          image_shape = box_preds.tensor.new_tensor(img_shape)
          valid_cam_inds = ((box_2d_preds[:, 0] < image_shape[1]) &
                            (box_2d_preds[:, 1] < image_shape[0]) &
                            (box_2d_preds[:, 2] > 0) & (box_2d_preds[:, 3] > 0))
          # check box_preds
          limit_range = box_preds.tensor.new_tensor(self.pcd_limit_range)
          valid_pcd_inds = ((box_preds.center > limit_range[:3]) &
                            (box_preds.center < limit_range[3:]))
          valid_inds = valid_cam_inds & valid_pcd_inds.all(-1)
  
          if valid_inds.sum() > 0:
              return dict(
                  bbox=box_2d_preds[valid_inds, :].numpy(),
                  box3d_camera=box_preds_camera[valid_inds].tensor.numpy(),
                  box3d_lidar=box_preds[valid_inds].tensor.numpy(),
                  scores=scores[valid_inds].numpy(),
                  customs_pred=customs_pred[valid_inds].numpy(),
                  label_preds=labels[valid_inds].numpy(),
                  sample_idx=sample_idx)
          else:
              return dict(
                  bbox=np.zeros([0, 4]),
                  box3d_camera=np.zeros([0, 7]),
                  box3d_lidar=np.zeros([0, 7]),
                  scores=np.zeros([0]),
                  customs_pred=np.zeros([0]),
                  label_preds=np.zeros([0, 4]),
                  sample_idx=sample_idx)
  
      def evaluate(self,
                   results,
                   metric=None,
                   logger=None,
                   pklfile_prefix=None,
                   submission_prefix=None,
                   show=False,
                   out_dir=None,
                   pipeline=None):
          """Evaluation in KITTI protocol.
  
          Args:
              results (list[dict]): Testing results of the dataset.
              metric (str | list[str]): Metrics to be evaluated.
              logger (logging.Logger | str | None): Logger used for printing
                  related information during evaluation. Default: None.
              pklfile_prefix (str | None): The prefix of pkl files. It includes
                  the file path and the prefix of filename, e.g., "a/b/prefix".
                  If not specified, a temp file will be created. Default: None.
              submission_prefix (str | None): The prefix of submission datas.
                  If not specified, the submission data will not be generated.
              show (bool): Whether to visualize.
                  Default: False.
              out_dir (str): Path to save the visualization results.
                  Default: None.
              pipeline (list[dict], optional): raw data loading for showing.
                  Default: None.
  
          Returns:
              dict[str, float]: Results of each evaluation metric.
          """
          result_files, tmp_dir = self.format_results(results, pklfile_prefix)
          from mmdet3d.core.evaluation import custom_kitti_eval
          gt_annos = [info['annos'] for info in self.data_infos]
          if self.custom_attr == 'distance':
              centers = [self.get_ann_info(idx)['gt_bboxes_3d'].tensor[:3] for idx in range(len(gt_annos))] # LiDAR coordinate
              gt_custom = [np.array(torch.norm(c, p=2, dim=1)) for c in centers]
              gt_custom = [np.pad(d,(0, len(gt_annos[idx]['name'])-len(d)), 'constant', constant_values=(0,-1)) for idx, d in enumerate(gt_custom)]
          else:
              gt_num_points = [info['gt_num_points'] for info in gt_annos]
              dims = [self.get_ann_info(idx)['gt_bboxes_3d'].tensor[3:6] for idx in range(len(gt_annos))] # LiDAR coordinate
              surf_area = [np.array((d[:, 0]+d[:, 1])*d[:, 2]) for d in dims]
              gt_custom = [num/s for num, s in zip(gt_num_points, surf_area)]
              gt_custom = [np.pad(d,(0, len(gt_annos[idx]['name'])-len(d)), 'constant', constant_values=(0,-1)) for idx, d in enumerate(gt_custom)]
          
          # change custom attribute to category
          import copy
          for i in range(len(gt_custom)):
              for label_d in range(len(self.custom_range)-1):
                  tmp = copy.deepcopy(gt_custom[i])
                  gt_custom[i][(self.custom_range[label_d]<tmp) & (tmp <= self.custom_range[label_d+1])] = label_d
          for i in range(len(result_files)):
              for label_d in range(len(self.custom_range)-1):
                  tmp = copy.deepcopy(result_files[i]['custom'])
                  result_files[i]['custom'][(self.custom_range[label_d]<tmp) & (tmp <= self.custom_range[label_d+1])] = label_d
          #_gt_custom = copy.deepcopy(gt_custom)
          for idx, gt in enumerate(gt_annos):
              gt.update({'custom': gt_custom[idx]})
          if isinstance(result_files, dict):
              ap_dict = dict()
              for name, result_files_ in result_files.items():
                  eval_types = ['bbox', 'bev', '3d']
                  if 'img' in name:
                      eval_types = ['bbox']
                  ap_result_str, ap_dict_ = custom_kitti_eval(
                      gt_annos,
                      result_files_,
                      self.CLASSES,
                      eval_types=eval_types,
                      num_custom=len(self.custom_range)-1)
                  for ap_type, ap in ap_dict_.items():
                      ap_dict[f'{name}/{ap_type}'] = float('{:.4f}'.format(ap))
  
                  print_log(
                      f'Results of {name}:\n' + ap_result_str, logger=logger)
  
          else:
              if metric == 'img_bbox':
                  ap_result_str, ap_dict = custom_kitti_eval(
                      gt_annos, result_files, self.CLASSES, eval_types=['bbox'], num_custom=len(self.custom_range)-1)
              else:
                  ap_result_str, ap_dict = custom_kitti_eval(gt_annos, result_files,
                                                      self.CLASSES, num_custom=len(self.custom_range)-1)
              print_log('\n' + ap_result_str, logger=logger)
  
          if tmp_dir is not None:
              tmp_dir.cleanup()
          if show:
              self.show(results, out_dir, pipeline=pipeline)
          return ap_dict
  
      def bbox2result_kitti(self,
                            net_outputs,
                            class_names,
                            pklfile_prefix=None,
                            submission_prefix=None):
          """Convert 3D detection results to kitti format for evaluation and test
          submission.
  
          Args:
              net_outputs (list[np.ndarray]): List of array storing the \
                  inferenced bounding boxes and scores.
              class_names (list[String]): A list of class names.
              pklfile_prefix (str | None): The prefix of pkl file.
              submission_prefix (str | None): The prefix of submission file.
  
          Returns:
              list[dict]: A list of dictionaries with the kitti format.
          """
          assert len(net_outputs) == len(self.data_infos), \
              'invalid list length of network outputs'
          if submission_prefix is not None:
              mmcv.mkdir_or_exist(submission_prefix)
  
          det_annos = []
          print('\nConverting prediction to KITTI format')
          for idx, pred_dicts in enumerate(
                  mmcv.track_iter_progress(net_outputs)):
              annos = []
              info = self.data_infos[idx]
              sample_idx = info['image']['image_idx']
              image_shape = info['image']['image_shape'][:2]
              box_dict = self.convert_valid_bboxes(pred_dicts, info)
              anno = {
                  'name': [],
                  'truncated': [],
                  'occluded': [],
                  'alpha': [],
                  'bbox': [],
                  'dimensions': [],
                  'location': [],
                  'rotation_y': [],
                  'score': [],
                  'custom_pred':[],
                  'custom':[]
              }
              if len(box_dict['bbox']) > 0:
                  box_2d_preds = box_dict['bbox']
                  box_preds = box_dict['box3d_camera']
                  scores = box_dict['scores']
                  customs_pred = box_dict['customs_pred']
                  box_preds_lidar = box_dict['box3d_lidar']
                  label_preds = box_dict['label_preds']
  
                  for box, box_lidar, bbox, score, custom_pred, label in zip(
                          box_preds, box_preds_lidar, box_2d_preds, scores, customs_pred,
                          label_preds):
                      bbox[2:] = np.minimum(bbox[2:], image_shape[::-1])
                      bbox[:2] = np.maximum(bbox[:2], [0, 0])
                      anno['name'].append(class_names[int(label)])
                      anno['truncated'].append(0.0)
                      anno['occluded'].append(0)
                      anno['alpha'].append(
                          -np.arctan2(-box_lidar[1], box_lidar[0]) + box[6])
                      anno['bbox'].append(bbox)
                      anno['dimensions'].append(box[3:6])
                      anno['location'].append(box[:3])
                      anno['rotation_y'].append(box[6])
                      anno['score'].append(score)
                      anno['custom_pred'].append(custom_pred),
                      custom = np.linalg.norm(box_lidar[0:3])
                      anno['custom'].append(custom)
  
                  anno = {k: np.stack(v) for k, v in anno.items()}
                  annos.append(anno)
              else:
                  anno = {
                      'name': np.array([]),
                      'truncated': np.array([]),
                      'occluded': np.array([]),
                      'alpha': np.array([]),
                      'bbox': np.zeros([0, 4]),
                      'dimensions': np.zeros([0, 3]),
                      'location': np.zeros([0, 3]),
                      'rotation_y': np.array([]),
                      'score': np.array([]),
                      'custom_pred': np.array([]),
                      'custom': np.array([]),
                  }
                  annos.append(anno)
  
              if submission_prefix is not None:
                  curr_file = f'{submission_prefix}/{sample_idx:06d}.txt'
                  with open(curr_file, 'w') as f:
                      bbox = anno['bbox']
                      loc = anno['location']
                      dims = anno['dimensions']  # lhw -> hwl
  
                      for idx in range(len(bbox)):
                          print(
                              '{} -1 -1 {:.4f} {:.4f} {:.4f} {:.4f} '
                              '{:.4f} {:.4f} {:.4f} '
                              '{:.4f} {:.4f} {:.4f} {:.4f} {:.4f} {:.4f}'.format(
                                  anno['name'][idx], anno['alpha'][idx],
                                  bbox[idx][0], bbox[idx][1], bbox[idx][2],
                                  bbox[idx][3], dims[idx][1], dims[idx][2],
                                  dims[idx][0], loc[idx][0], loc[idx][1],
                                  loc[idx][2], anno['rotation_y'][idx],
                                  anno['score'][idx]),
                              file=f)
  
              annos[-1]['sample_idx'] = np.array(
                  [sample_idx] * len(annos[-1]['score']), dtype=np.int64)
  
              det_annos += annos
  
          if pklfile_prefix is not None:
              if not pklfile_prefix.endswith(('.pkl', '.pickle')):
                  out = f'{pklfile_prefix}.pkl'
              mmcv.dump(det_annos, out)
              print(f'Result is saved to {out}.')
  
          return det_annos
  