mirror of
https://github.com/alibaba/EasyCV.git
synced 2025-06-03 14:49:00 +08:00
322 lines
12 KiB
Python
322 lines
12 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import os
|
|
|
|
import mmcv
|
|
import numpy as np
|
|
|
|
from easycv.core.bbox import LiDARInstance3DBoxes
|
|
from easycv.datasets.registry import DATASOURCES
|
|
from .base import Det3dSourceBase
|
|
|
|
|
|
@DATASOURCES.register_module()
|
|
class Det3dSourceNuScenes(Det3dSourceBase):
|
|
r"""NuScenes data source.
|
|
|
|
This class serves as the API for experiments on the NuScenes Dataset.
|
|
|
|
Please refer to `NuScenes Dataset <https://www.nuscenes.org/download>`_
|
|
for data downloading.
|
|
|
|
Args:
|
|
ann_file (str): Path of annotation file.
|
|
pipeline (list[dict], optional): Pipeline used for data processing.
|
|
Defaults to None.
|
|
data_root (str): Path of dataset root.
|
|
classes (tuple[str], optional): Classes used in the dataset.
|
|
Defaults to None.
|
|
load_interval (int, optional): Interval of loading the dataset. It is
|
|
used to uniformly sample the dataset. Defaults to 1.
|
|
with_velocity (bool, optional): Whether include velocity prediction
|
|
into the experiments. Defaults to True.
|
|
modality (dict, optional): Modality to specify the sensor data used
|
|
as input. Defaults to None.
|
|
box_type_3d (str, optional): Type of 3D box of this dataset.
|
|
Based on the `box_type_3d`, the dataset will encapsulate the box
|
|
to its original format then converted them to `box_type_3d`.
|
|
Defaults to 'LiDAR' in this dataset. Available options includes.
|
|
- 'LiDAR': Box in LiDAR coordinates.
|
|
- 'Depth': Box in depth coordinates, usually for indoor dataset.
|
|
- 'Camera': Box in camera coordinates.
|
|
filter_empty_gt (bool, optional): Whether to filter empty GT.
|
|
Defaults to True.
|
|
test_mode (bool, optional): Whether the dataset is in test mode.
|
|
Defaults to False.
|
|
use_valid_flag (bool, optional): Whether to use `use_valid_flag` key
|
|
in the info file as mask to filter gt_boxes and gt_names.
|
|
Defaults to False.
|
|
"""
|
|
NameMapping = {
|
|
'movable_object.barrier': 'barrier',
|
|
'vehicle.bicycle': 'bicycle',
|
|
'vehicle.bus.bendy': 'bus',
|
|
'vehicle.bus.rigid': 'bus',
|
|
'vehicle.car': 'car',
|
|
'vehicle.construction': 'construction_vehicle',
|
|
'vehicle.motorcycle': 'motorcycle',
|
|
'human.pedestrian.adult': 'pedestrian',
|
|
'human.pedestrian.child': 'pedestrian',
|
|
'human.pedestrian.construction_worker': 'pedestrian',
|
|
'human.pedestrian.police_officer': 'pedestrian',
|
|
'movable_object.trafficcone': 'traffic_cone',
|
|
'vehicle.trailer': 'trailer',
|
|
'vehicle.truck': 'truck'
|
|
}
|
|
DefaultAttribute = {
|
|
'car': 'vehicle.parked',
|
|
'pedestrian': 'pedestrian.moving',
|
|
'trailer': 'vehicle.parked',
|
|
'truck': 'vehicle.parked',
|
|
'bus': 'vehicle.moving',
|
|
'motorcycle': 'cycle.without_rider',
|
|
'construction_vehicle': 'vehicle.parked',
|
|
'bicycle': 'cycle.without_rider',
|
|
'barrier': '',
|
|
'traffic_cone': '',
|
|
}
|
|
AttrMapping = {
|
|
'cycle.with_rider': 0,
|
|
'cycle.without_rider': 1,
|
|
'pedestrian.moving': 2,
|
|
'pedestrian.standing': 3,
|
|
'pedestrian.sitting_lying_down': 4,
|
|
'vehicle.moving': 5,
|
|
'vehicle.parked': 6,
|
|
'vehicle.stopped': 7,
|
|
}
|
|
AttrMapping_rev = [
|
|
'cycle.with_rider',
|
|
'cycle.without_rider',
|
|
'pedestrian.moving',
|
|
'pedestrian.standing',
|
|
'pedestrian.sitting_lying_down',
|
|
'vehicle.moving',
|
|
'vehicle.parked',
|
|
'vehicle.stopped',
|
|
]
|
|
# https://github.com/nutonomy/nuscenes-devkit/blob/57889ff20678577025326cfc24e57424a829be0a/python-sdk/nuscenes/eval/detection/evaluate.py#L222 # noqa
|
|
ErrNameMapping = {
|
|
'trans_err': 'mATE',
|
|
'scale_err': 'mASE',
|
|
'orient_err': 'mAOE',
|
|
'vel_err': 'mAVE',
|
|
'attr_err': 'mAAE'
|
|
}
|
|
CLASSES = ('car', 'truck', 'trailer', 'bus', 'construction_vehicle',
|
|
'bicycle', 'motorcycle', 'pedestrian', 'traffic_cone',
|
|
'barrier')
|
|
|
|
def __init__(
|
|
self,
|
|
ann_file,
|
|
pipeline=None,
|
|
data_root=None,
|
|
classes=None,
|
|
load_interval=1,
|
|
with_velocity=True,
|
|
modality=None,
|
|
box_type_3d='LiDAR',
|
|
filter_empty_gt=True,
|
|
test_mode=False,
|
|
use_valid_flag=False,
|
|
):
|
|
self.load_interval = load_interval
|
|
self.use_valid_flag = use_valid_flag
|
|
self.with_velocity = with_velocity
|
|
|
|
super().__init__(
|
|
data_root=data_root,
|
|
ann_file=ann_file,
|
|
pipeline=pipeline,
|
|
classes=classes,
|
|
modality=modality,
|
|
box_type_3d=box_type_3d,
|
|
filter_empty_gt=filter_empty_gt,
|
|
test_mode=test_mode)
|
|
|
|
if self.modality is None:
|
|
self.modality = dict(
|
|
use_camera=False,
|
|
use_lidar=True,
|
|
use_radar=False,
|
|
use_map=False,
|
|
use_external=False,
|
|
)
|
|
|
|
def get_cat_ids(self, idx):
|
|
"""Get category distribution of single scene.
|
|
|
|
Args:
|
|
idx (int): Index of the data_info.
|
|
|
|
Returns:
|
|
dict[list]: for each category, if the current scene
|
|
contains such boxes, store a list containing idx,
|
|
otherwise, store empty list.
|
|
"""
|
|
info = self.data_infos[idx]
|
|
if self.use_valid_flag:
|
|
mask = info['valid_flag']
|
|
gt_names = set(info['gt_names'][mask])
|
|
else:
|
|
gt_names = set(info['gt_names'])
|
|
|
|
cat_ids = []
|
|
for name in gt_names:
|
|
if name in self.CLASSES:
|
|
cat_ids.append(self.cat2id[name])
|
|
return cat_ids
|
|
|
|
def load_annotations(self, ann_file):
|
|
"""Load annotations from ann_file.
|
|
|
|
Args:
|
|
ann_file (str): Path of the annotation file.
|
|
|
|
Returns:
|
|
list[dict]: List of annotations sorted by timestamps.
|
|
"""
|
|
data = mmcv.load(ann_file, file_format='pkl')
|
|
data_infos = list(sorted(data['infos'], key=lambda e: e['timestamp']))
|
|
data_infos = data_infos[::self.load_interval]
|
|
self.metadata = data['metadata']
|
|
self.version = self.metadata['version']
|
|
return data_infos
|
|
|
|
def get_data_info(self, index):
|
|
"""Get data info according to the given index.
|
|
|
|
Args:
|
|
index (int): Index of the sample data to get.
|
|
|
|
Returns:
|
|
dict: Data information that will be passed to the data \
|
|
preprocessing pipelines. It includes the following keys:
|
|
|
|
- sample_idx (str): Sample index.
|
|
- pts_filename (str): Filename of point clouds.
|
|
- sweeps (list[dict]): Infos of sweeps.
|
|
- timestamp (float): Sample timestamp.
|
|
- img_filename (str, optional): Image filename.
|
|
- lidar2img (list[np.ndarray], optional): Transformations \
|
|
from lidar to different cameras.
|
|
- ann_info (dict): Annotation info.
|
|
"""
|
|
from nuscenes.eval.common.utils import Quaternion, quaternion_yaw
|
|
|
|
info = self.data_infos[index]
|
|
# standard protocal modified from SECOND.Pytorch
|
|
input_dict = dict(
|
|
sample_idx=info['token'],
|
|
pts_filename=info['lidar_path'],
|
|
sweeps=info['sweeps'],
|
|
ego2global_translation=info['ego2global_translation'],
|
|
ego2global_rotation=info['ego2global_rotation'],
|
|
prev_idx=info['prev'],
|
|
next_idx=info['next'],
|
|
scene_token=info['scene_token'],
|
|
can_bus=info['can_bus'],
|
|
frame_idx=info['frame_idx'],
|
|
timestamp=info['timestamp'] / 1e6,
|
|
)
|
|
|
|
if self.modality['use_camera']:
|
|
image_paths = []
|
|
lidar2img_rts = []
|
|
lidar2cam_rts = []
|
|
cam_intrinsics = []
|
|
for cam_type, cam_info in info['cams'].items():
|
|
cam_info['data_path'] = os.path.join(self.data_root,
|
|
cam_info['data_path'])
|
|
image_paths.append(cam_info['data_path'])
|
|
# obtain lidar to image transformation matrix
|
|
lidar2cam_r = np.linalg.inv(cam_info['sensor2lidar_rotation'])
|
|
lidar2cam_t = cam_info[
|
|
'sensor2lidar_translation'] @ lidar2cam_r.T
|
|
lidar2cam_rt = np.eye(4)
|
|
lidar2cam_rt[:3, :3] = lidar2cam_r.T
|
|
lidar2cam_rt[3, :3] = -lidar2cam_t
|
|
intrinsic = cam_info['cam_intrinsic']
|
|
viewpad = np.eye(4)
|
|
viewpad[:intrinsic.shape[0], :intrinsic.shape[1]] = intrinsic
|
|
lidar2img_rt = (viewpad @ lidar2cam_rt.T)
|
|
lidar2img_rts.append(lidar2img_rt)
|
|
|
|
cam_intrinsics.append(viewpad)
|
|
lidar2cam_rts.append(lidar2cam_rt.T)
|
|
|
|
input_dict.update(
|
|
dict(
|
|
img_filename=image_paths,
|
|
lidar2img=lidar2img_rts,
|
|
cam_intrinsic=cam_intrinsics,
|
|
lidar2cam=lidar2cam_rts,
|
|
))
|
|
|
|
if not self.test_mode:
|
|
annos = self.get_ann_info(index)
|
|
input_dict['ann_info'] = annos
|
|
|
|
rotation = Quaternion(input_dict['ego2global_rotation'])
|
|
translation = input_dict['ego2global_translation']
|
|
can_bus = input_dict['can_bus']
|
|
can_bus[:3] = translation
|
|
can_bus[3:7] = rotation
|
|
patch_angle = quaternion_yaw(rotation) / np.pi * 180
|
|
if patch_angle < 0:
|
|
patch_angle += 360
|
|
can_bus[-2] = patch_angle / 180 * np.pi
|
|
can_bus[-1] = patch_angle
|
|
|
|
return input_dict
|
|
|
|
def get_ann_info(self, index):
|
|
"""Get annotation info according to the given index.
|
|
|
|
Args:
|
|
index (int): Index of the annotation data to get.
|
|
|
|
Returns:
|
|
dict: Annotation information consists of the following keys:
|
|
|
|
- gt_bboxes_3d (:obj:`LiDARInstance3DBoxes`):
|
|
3D ground truth bboxes
|
|
- gt_labels_3d (np.ndarray): Labels of ground truths.
|
|
- gt_names (list[str]): Class names of ground truths.
|
|
"""
|
|
info = self.data_infos[index]
|
|
# filter out bbox containing no points
|
|
if self.use_valid_flag:
|
|
mask = info['valid_flag']
|
|
else:
|
|
mask = info['num_lidar_pts'] > 0
|
|
gt_bboxes_3d = info['gt_boxes'][mask]
|
|
gt_names_3d = info['gt_names'][mask]
|
|
gt_labels_3d = []
|
|
for cat in gt_names_3d:
|
|
if cat in self.CLASSES:
|
|
gt_labels_3d.append(self.CLASSES.index(cat))
|
|
else:
|
|
gt_labels_3d.append(-1)
|
|
gt_labels_3d = np.array(gt_labels_3d)
|
|
|
|
if self.with_velocity:
|
|
gt_velocity = info['gt_velocity'][mask]
|
|
nan_mask = np.isnan(gt_velocity[:, 0])
|
|
gt_velocity[nan_mask] = [0.0, 0.0]
|
|
gt_bboxes_3d = np.concatenate([gt_bboxes_3d, gt_velocity], axis=-1)
|
|
|
|
# the nuscenes box center is [0.5, 0.5, 0.5], we change it to be
|
|
# the same as KITTI (0.5, 0.5, 0)
|
|
gt_bboxes_3d = LiDARInstance3DBoxes(
|
|
gt_bboxes_3d,
|
|
box_dim=gt_bboxes_3d.shape[-1],
|
|
origin=(0.5, 0.5, 0.5)).convert_to(self.box_mode_3d)
|
|
|
|
anns_results = dict(
|
|
gt_bboxes_3d=gt_bboxes_3d,
|
|
gt_labels_3d=gt_labels_3d,
|
|
gt_names=gt_names_3d)
|
|
return anns_results
|