Skip to content

GitLab

Commit 71b88ff5 authored by jshilong's avatar jshilong Committed by ChaimZhu
Browse files

Refactor the Scannet and SunRGBD

parent 39b294f5
Hide whitespace changes
Inline Side-by-side
Showing with 471 additions and 286 deletions
mmdet3d/datasets/det3d_dataset.py
View file @ 71b88ff5
......@@ -83,7 +83,7 @@ class Det3DDataset(BaseDataset):
self.modality = modality
assert self.modality['use_lidar'] or self.modality['use_camera'], (
'Please specify the `modality` (`use_lidar` '
f' or `use_camera`) for {self.__class__.__name__}')
f', `use_camera`) for {self.__class__.__name__}')
self.box_type_3d, self.box_mode_3d = get_box_type(box_type_3d)
......
mmdet3d/datasets/kitti_dataset.py
View file @ 71b88ff5
......@@ -31,7 +31,7 @@ class KittiDataset(Det3DDataset):
pipeline (list[dict], optional): Pipeline used for data processing.
Defaults to None.
modality (dict, optional): Modality to specify the sensor data used
as input. Defaults to None.
as input. Defaults to `dict(use_lidar=True)`.
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`.
......@@ -55,7 +55,7 @@ class KittiDataset(Det3DDataset):
data_root: str,
ann_file: str,
pipeline: List[Union[dict, Callable]] = [],
modality: Optional[dict] = None,
modality: Optional[dict] = dict(use_lidar=True),
box_type_3d: str = 'LiDAR',
filter_empty_gt: bool = True,
test_mode: bool = False,
......
mmdet3d/datasets/scannet_dataset.py
View file @ 71b88ff5
......@@ -2,6 +2,7 @@
import tempfile
import warnings
from os import path as osp
from typing import Callable, List, Union
import numpy as np
......@@ -25,13 +26,17 @@ class ScanNetDataset(Det3DDataset):
Args:
data_root (str): Path of dataset root.
ann_file (str): Path of annotation file.
pipeline (list[dict], optional): Pipeline used for data processing.
Defaults to None.
classes (tuple[str], optional): Classes used in the dataset.
metainfo (dict, optional): Meta information for dataset, such as class
information. Defaults to None.
data_prefix (dict): Prefix for data. Defaults to
`dict(pts='points',
pts_isntance_mask='instance_mask',
pts_semantic_mask='semantic_mask')`.
pipeline (list[dict]): Pipeline used for data processing.
Defaults to None.
modality (dict, optional): Modality to specify the sensor data used
modality (dict): Modality to specify the sensor data used
as input. Defaults to None.
box_type_3d (str, optional): Type of 3D box of this dataset.
box_type_3d (str): 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 'Depth' in this dataset. Available options includes
......@@ -39,176 +44,113 @@ class ScanNetDataset(Det3DDataset):
- '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.
filter_empty_gt (bool): Whether to filter empty GT.
Defaults to True.
test_mode (bool, optional): Whether the dataset is in test mode.
test_mode (bool): Whether the dataset is in test mode.
Defaults to False.
"""
CLASSES = ('cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window',
'bookshelf', 'picture', 'counter', 'desk', 'curtain',
'refrigerator', 'showercurtrain', 'toilet', 'sink', 'bathtub',
'garbagebin')
METAINFO = {
'CLASSES':
('cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window',
'bookshelf', 'picture', 'counter', 'desk', 'curtain', 'refrigerator',
'showercurtrain', 'toilet', 'sink', 'bathtub', 'garbagebin')
}
def __init__(self,
data_root,
ann_file,
pipeline=None,
classes=None,
modality=dict(use_camera=False, use_depth=True),
box_type_3d='Depth',
filter_empty_gt=True,
test_mode=False,
data_root: str,
ann_file: str,
metainfo: dict = None,
data_prefix: dict = dict(
pts='points',
pts_isntance_mask='instance_mask',
pts_semantic_mask='semantic_mask'),
pipeline: List[Union[dict, Callable]] = [],
modality=dict(use_camera=False, use_lidar=True),
box_type_3d: str = 'Depth',
filter_empty_gt: bool = True,
test_mode: bool = False,
**kwargs):
super().__init__(
data_root=data_root,
ann_file=ann_file,
metainfo=metainfo,
data_prefix=data_prefix,
pipeline=pipeline,
classes=classes,
modality=modality,
box_type_3d=box_type_3d,
filter_empty_gt=filter_empty_gt,
test_mode=test_mode,
**kwargs)
assert 'use_camera' in self.modality and \
'use_depth' in self.modality
assert self.modality['use_camera'] or self.modality['use_depth']
def get_data_info(self, index):
"""Get data info according to the given index.
Args:
index (int): Index of the sample data to get.
'use_lidar' in self.modality
assert self.modality['use_camera'] or self.modality['use_lidar']
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.
- file_name (str): Filename of point clouds.
- img_prefix (str, optional): Prefix of image files.
- img_info (dict, optional): Image info.
- ann_info (dict): Annotation info.
"""
info = self.data_infos[index]
sample_idx = info['point_cloud']['lidar_idx']
pts_filename = osp.join(self.data_root, info['pts_path'])
input_dict = dict(sample_idx=sample_idx)
if self.modality['use_depth']:
input_dict['pts_filename'] = pts_filename
input_dict['file_name'] = pts_filename
if self.modality['use_camera']:
img_info = []
for img_path in info['img_paths']:
img_info.append(
dict(filename=osp.join(self.data_root, img_path)))
intrinsic = info['intrinsics']
axis_align_matrix = self._get_axis_align_matrix(info)
depth2img = []
for extrinsic in info['extrinsics']:
depth2img.append(
intrinsic @ np.linalg.inv(axis_align_matrix @ extrinsic))
input_dict['img_prefix'] = None
input_dict['img_info'] = img_info
input_dict['depth2img'] = depth2img
if not self.test_mode:
annos = self.get_ann_info(index)
input_dict['ann_info'] = annos
if self.filter_empty_gt and ~(annos['gt_labels_3d'] != -1).any():
return None
return input_dict
def get_ann_info(self, index):
"""Get annotation info according to the given index.
@staticmethod
def _get_axis_align_matrix(info: dict) -> dict:
"""Get axis_align_matrix from info. If not exist, return identity mat.
Args:
index (int): Index of the annotation data to get.
info (dict): Info of a single sample data.
Returns:
dict: annotation information consists of the following keys:
- gt_bboxes_3d (:obj:`DepthInstance3DBoxes`):
3D ground truth bboxes
- gt_labels_3d (np.ndarray): Labels of ground truths.
- pts_instance_mask_path (str): Path of instance masks.
- pts_semantic_mask_path (str): Path of semantic masks.
- axis_align_matrix (np.ndarray): Transformation matrix for
global scene alignment.
np.ndarray: 4x4 transformation matrix.
"""
# Use index to get the annos, thus the evalhook could also use this api
info = self.data_infos[index]
if info['annos']['gt_num'] != 0:
gt_bboxes_3d = info['annos']['gt_boxes_upright_depth'].astype(
np.float32) # k, 6
gt_labels_3d = info['annos']['class'].astype(np.int64)
if 'axis_align_matrix' in info:
return np.array(info['axis_align_matrix'])
else:
gt_bboxes_3d = np.zeros((0, 6), dtype=np.float32)
gt_labels_3d = np.zeros((0, ), dtype=np.int64)
# to target box structure
gt_bboxes_3d = DepthInstance3DBoxes(
gt_bboxes_3d,
box_dim=gt_bboxes_3d.shape[-1],
with_yaw=False,
origin=(0.5, 0.5, 0.5)).convert_to(self.box_mode_3d)
pts_instance_mask_path = osp.join(self.data_root,
info['pts_instance_mask_path'])
pts_semantic_mask_path = osp.join(self.data_root,
info['pts_semantic_mask_path'])
axis_align_matrix = self._get_axis_align_matrix(info)
anns_results = dict(
gt_bboxes_3d=gt_bboxes_3d,
gt_labels_3d=gt_labels_3d,
pts_instance_mask_path=pts_instance_mask_path,
pts_semantic_mask_path=pts_semantic_mask_path,
axis_align_matrix=axis_align_matrix)
return anns_results
warnings.warn(
'axis_align_matrix is not found in ScanNet data info, please '
'use new pre-process scripts to re-generate ScanNet data')
return np.eye(4).astype(np.float32)
def prepare_test_data(self, index):
"""Prepare data for testing.
def parse_data_info(self, info: dict) -> dict:
"""Process the raw data info.
We should take axis_align_matrix from self.data_infos since we need
to align point clouds.
The only difference with it in `Det3DDataset`
is the specific process for `axis_align_matrix'.
Args:
index (int): Index for accessing the target data.
info (dict): Raw info dict.
Returns:
dict: Testing data dict of the corresponding index.
dict: Data information that will be passed to the data
preprocessing pipelines. It includes the following keys:
"""
input_dict = self.get_data_info(index)
# take the axis_align_matrix from data_infos
input_dict['ann_info'] = dict(
axis_align_matrix=self._get_axis_align_matrix(
self.data_infos[index]))
self.pre_pipeline(input_dict)
example = self.pipeline(input_dict)
return example
@staticmethod
def _get_axis_align_matrix(info):
"""Get axis_align_matrix from info. If not exist, return identity mat.
# TODO: whether all depth modality is pts ?
if self.modality['use_lidar']:
info['lidar_points']['lidar_path'] = \
osp.join(
self.data_prefix.get('pts', ''),
info['lidar_points']['lidar_path'])
info['axis_align_matrix'] = self._get_axis_align_matrix(info)
info['pts_instance_mask_path'] = osp.join(
self.data_prefix.get('pts_instance_mask', ''),
info['pts_instance_mask_path'])
info['pts_semantic_mask_path'] = osp.join(
self.data_prefix.get('pts_semantic_mask', ''),
info['pts_semantic_mask_path'])
info = super().parse_data_info(info)
return info
def parse_ann_info(self, info: dict) -> dict:
"""Process the `instances` in data info to `ann_info`
Args:
info (dict): one data info term.
info (dict): Info dict.
Returns:
np.ndarray: 4x4 transformation matrix.
dict: Processed `ann_info`
"""
if 'axis_align_matrix' in info['annos'].keys():
return info['annos']['axis_align_matrix'].astype(np.float32)
else:
warnings.warn(
'axis_align_matrix is not found in ScanNet data info, please '
'use new pre-process scripts to re-generate ScanNet data')
return np.eye(4).astype(np.float32)
ann_info = super().parse_ann_info(info)
# to target box structure
ann_info['gt_bboxes_3d'] = DepthInstance3DBoxes(
ann_info['gt_bboxes_3d'],
box_dim=ann_info['gt_bboxes_3d'].shape[-1],
with_yaw=False,
origin=(0.5, 0.5, 0.5)).convert_to(self.box_mode_3d)
return ann_info
def _build_default_pipeline(self):
"""Build the default pipeline for this dataset."""
......@@ -241,8 +183,8 @@ class ScanNetDataset(Det3DDataset):
assert out_dir is not None, 'Expect out_dir, got none.'
pipeline = self._get_pipeline(pipeline)
for i, result in enumerate(results):
data_info = self.data_infos[i]
pts_path = data_info['pts_path']
data_info = self.get_data_info[i]
pts_path = data_info['lidar_points']['lidar_path']
file_name = osp.split(pts_path)[-1].split('.')[0]
points = self._extract_data(i, pipeline, 'points').numpy()
gt_bboxes = self.get_ann_info(i)['gt_bboxes_3d'].tensor.numpy()
......
mmdet3d/datasets/sunrgbd_dataset.py
View file @ 71b88ff5
# Copyright (c) OpenMMLab. All rights reserved.
from collections import OrderedDict
from os import path as osp
import numpy as np
from typing import Callable, List, Optional, Union
from mmdet3d.core import show_multi_modality_result, show_result
from mmdet3d.core.bbox import DepthInstance3DBoxes
......@@ -24,12 +23,14 @@ class SUNRGBDDataset(Det3DDataset):
Args:
data_root (str): Path of dataset root.
ann_file (str): Path of annotation file.
metainfo (dict, optional): Meta information for dataset, such as class
information. Defaults to None.
data_prefix (dict, optional): Prefix for data. Defaults to
`dict(pts='points',img='sunrgbd_trainval')`.
pipeline (list[dict], optional): Pipeline used for data processing.
Defaults to None.
classes (tuple[str], optional): Classes used in the dataset.
Defaults to None.
modality (dict, optional): Modality to specify the sensor data used
as input. Defaults to None.
as input. Defaults to `dict(use_camera=True, use_lidar=True)`.
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`.
......@@ -43,24 +44,29 @@ class SUNRGBDDataset(Det3DDataset):
test_mode (bool, optional): Whether the dataset is in test mode.
Defaults to False.
"""
CLASSES = ('bed', 'table', 'sofa', 'chair', 'toilet', 'desk', 'dresser',
'night_stand', 'bookshelf', 'bathtub')
METAINFO = {
'CLASSES': ('bed', 'table', 'sofa', 'chair', 'toilet', 'desk',
'dresser', 'night_stand', 'bookshelf', 'bathtub')
}
def __init__(self,
data_root,
ann_file,
pipeline=None,
classes=None,
data_root: str,
ann_file: str,
metainfo: Optional[dict] = None,
data_prefix: dict = dict(
pts='points', img='sunrgbd_trainval'),
pipeline: List[Union[dict, Callable]] = [],
modality=dict(use_camera=True, use_lidar=True),
box_type_3d='Depth',
filter_empty_gt=True,
test_mode=False,
box_type_3d: str = 'Depth',
filter_empty_gt: bool = True,
test_mode: bool = False,
**kwargs):
super().__init__(
data_root=data_root,
ann_file=ann_file,
metainfo=metainfo,
data_prefix=data_prefix,
pipeline=pipeline,
classes=classes,
modality=modality,
box_type_3d=box_type_3d,
filter_empty_gt=filter_empty_gt,
......@@ -70,96 +76,22 @@ class SUNRGBDDataset(Det3DDataset):
'use_lidar' in self.modality
assert self.modality['use_camera'] or self.modality['use_lidar']
def get_data_info(self, index):
"""Get data info according to the given index.
def parse_ann_info(self, info: dict) -> dict:
"""Process the `instances` in data info to `ann_info`
Args:
index (int): Index of the sample data to get.
info (dict): Info dict.
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, optional): Filename of point clouds.
- file_name (str, optional): Filename of point clouds.
- img_prefix (str, optional): Prefix of image files.
- img_info (dict, optional): Image info.
- calib (dict, optional): Camera calibration info.
- ann_info (dict): Annotation info.
dict: Processed `ann_info`
"""
info = self.data_infos[index]
sample_idx = info['point_cloud']['lidar_idx']
assert info['point_cloud']['lidar_idx'] == info['image']['image_idx']
input_dict = dict(sample_idx=sample_idx)
if self.modality['use_lidar']:
pts_filename = osp.join(self.data_root, info['pts_path'])
input_dict['pts_filename'] = pts_filename
input_dict['file_name'] = pts_filename
if self.modality['use_camera']:
img_filename = osp.join(
osp.join(self.data_root, 'sunrgbd_trainval'),
info['image']['image_path'])
input_dict['img_prefix'] = None
input_dict['img_info'] = dict(filename=img_filename)
calib = info['calib']
rt_mat = calib['Rt']
# follow Coord3DMode.convert_point
rt_mat = np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]
]) @ rt_mat.transpose(1, 0)
depth2img = calib['K'] @ rt_mat
input_dict['depth2img'] = depth2img
if not self.test_mode:
annos = self.get_ann_info(index)
input_dict['ann_info'] = annos
if self.filter_empty_gt and len(annos['gt_bboxes_3d']) == 0:
return None
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:`DepthInstance3DBoxes`):
3D ground truth bboxes
- gt_labels_3d (np.ndarray): Labels of ground truths.
- pts_instance_mask_path (str): Path of instance masks.
- pts_semantic_mask_path (str): Path of semantic masks.
"""
# Use index to get the annos, thus the evalhook could also use this api
info = self.data_infos[index]
if info['annos']['gt_num'] != 0:
gt_bboxes_3d = info['annos']['gt_boxes_upright_depth'].astype(
np.float32) # k, 6
gt_labels_3d = info['annos']['class'].astype(np.int64)
else:
gt_bboxes_3d = np.zeros((0, 7), dtype=np.float32)
gt_labels_3d = np.zeros((0, ), dtype=np.int64)
ann_info = super().parse_ann_info(info)
# to target box structure
gt_bboxes_3d = DepthInstance3DBoxes(
gt_bboxes_3d, 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)
if self.modality['use_camera']:
if info['annos']['gt_num'] != 0:
gt_bboxes_2d = info['annos']['bbox'].astype(np.float32)
else:
gt_bboxes_2d = np.zeros((0, 4), dtype=np.float32)
anns_results['bboxes'] = gt_bboxes_2d
anns_results['labels'] = gt_labels_3d
ann_info['gt_bboxes_3d'] = DepthInstance3DBoxes(
ann_info['gt_bboxes_3d'],
origin=(0.5, 0.5, 0.5)).convert_to(self.box_mode_3d)
return anns_results
return ann_info
def _build_default_pipeline(self):
"""Build the default pipeline for this dataset."""
......@@ -180,6 +112,7 @@ class SUNRGBDDataset(Det3DDataset):
pipeline.insert(0, dict(type='LoadImageFromFile'))
return Compose(pipeline)
# TODO fix this
def show(self, results, out_dir, show=True, pipeline=None):
"""Results visualization.
......
tests/test_data/test_datasets/test_kitti_dataset.py
View file @ 71b88ff5
# Copyright (c) OpenMMLab. All rights reserved.
import numpy as np
import torch
from mmcv.transforms.base import BaseTransform
from mmengine.registry import TRANSFORMS
from mmdet3d.core import LiDARInstance3DBoxes
from mmdet3d.datasets import KittiDataset
......@@ -11,35 +14,21 @@ def _generate_kitti_dataset_config():
ann_file = 'kitti_infos_train.pkl'
classes = ['Pedestrian', 'Cyclist', 'Car']
# wait for pipline refactor
if 'Identity' not in TRANSFORMS:
@TRANSFORMS.register_module()
class Identity(BaseTransform):
def transform(self, info):
if 'ann_info' in info:
info['gt_labels_3d'] = info['ann_info']['gt_labels_3d']
return info
pipeline = [
dict(
type='LoadPointsFromFile',
coord_type='LIDAR',
load_dim=4,
use_dim=4,
file_client_args=dict(backend='disk')),
dict(
type='MultiScaleFlipAug3D',
img_scale=(1333, 800),
pts_scale_ratio=1,
flip=False,
transforms=[
dict(
type='GlobalRotScaleTrans',
rot_range=[0, 0],
scale_ratio_range=[1.0, 1.0],
translation_std=[0, 0, 0]),
dict(type='RandomFlip3D'),
dict(
type='PointsRangeFilter',
point_cloud_range=[0, -40, -3, 70.4, 40, 1]),
dict(
type='DefaultFormatBundle3D',
class_names=classes,
with_label=False),
dict(type='Collect3D', keys=['points'])
])
dict(type='Identity'),
]
modality = dict(use_lidar=True, use_camera=False)
data_prefix = dict(pts='training/velodyne_reduced', img='training/image_2')
return data_root, ann_file, classes, data_prefix, pipeline, modality
......@@ -48,23 +37,9 @@ def _generate_kitti_dataset_config():
def test_getitem():
np.random.seed(0)
data_root, ann_file, classes, data_prefix, \
_, modality, = _generate_kitti_dataset_config()
pipeline, modality, = _generate_kitti_dataset_config()
modality['use_camera'] = True
from mmcv.transforms.base import BaseTransform
from mmengine.registry import TRANSFORMS
@TRANSFORMS.register_module()
class Identity(BaseTransform):
def transform(self, info):
if 'ann_info' in info:
info['gt_labels_3d'] = info['ann_info']['gt_labels_3d']
return info
pipeline = [
dict(type='Identity'),
]
kitti_dataset = KittiDataset(
data_root,
ann_file,
......@@ -94,12 +69,15 @@ def test_getitem():
assert ann_info['gt_labels'].dtype == np.int64
# only one instance
assert len(ann_info['gt_labels']) == 1
assert (ann_info['gt_labels'] == 0).all()
assert 'gt_labels_3d' in ann_info
assert ann_info['gt_labels_3d'].dtype == np.int64
assert 'gt_bboxes' in ann_info
assert ann_info['gt_bboxes'].dtype == np.float64
assert 'gt_bboxes_3d' in ann_info
assert isinstance(ann_info['gt_bboxes_3d'], LiDARInstance3DBoxes)
assert torch.allclose(ann_info['gt_bboxes_3d'].tensor.sum(),
torch.tensor(7.2650))
assert 'group_id' in ann_info
assert ann_info['group_id'].dtype == np.int64
assert 'occluded' in ann_info
......
tests/test_data/test_datasets/test_scannet_dataset.py 0 → 100644
View file @ 71b88ff5
# Copyright (c) OpenMMLab. All rights reserved.
import unittest
import numpy as np
import torch
from mmengine.testing import assert_allclose
from mmdet3d.core import DepthInstance3DBoxes
from mmdet3d.datasets import ScanNetDataset
def _generate_scannet_dataset_config():
data_root = 'tests/data/scannet'
ann_file = 'scannet_infos.pkl'
classes = ('cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window',
'bookshelf', 'picture', 'counter', 'desk', 'curtain',
'refrigerator', 'showercurtrain', 'toilet', 'sink', 'bathtub',
'garbagebin')
# TODO add pipline
from mmcv.transforms.base import BaseTransform
from mmengine.registry import TRANSFORMS
if 'Identity' not in TRANSFORMS:
@TRANSFORMS.register_module()
class Identity(BaseTransform):
def transform(self, info):
if 'ann_info' in info:
info['gt_labels_3d'] = info['ann_info']['gt_labels_3d']
return info
modality = dict(use_lidar=True, use_camera=False)
pipeline = [
dict(type='Identity'),
]
data_prefix = dict(
pts='points',
pts_instance_mask='instance_mask',
pts_semantic_mask='semantic_mask')
return data_root, ann_file, classes, data_prefix, pipeline, modality
class TestScanNetDataset(unittest.TestCase):
def test_scannet(self):
np.random.seed(0)
data_root, ann_file, classes, data_prefix, \
pipeline, modality, = _generate_scannet_dataset_config()
scannet_dataset = ScanNetDataset(
data_root,
ann_file,
data_prefix=data_prefix,
pipeline=pipeline,
metainfo=dict(CLASSES=classes),
modality=modality)
scannet_dataset.prepare_data(0)
input_dict = scannet_dataset.get_data_info(0)
scannet_dataset[0]
# assert the the path should contains data_prefix and data_root
self.assertIn(data_prefix['pts'],
input_dict['lidar_points']['lidar_path'])
self.assertIn(data_root, input_dict['lidar_points']['lidar_path'])
ann_info = scannet_dataset.parse_ann_info(input_dict)
# assert the keys in ann_info and the type
except_label = np.array([
6, 6, 4, 9, 11, 11, 10, 0, 15, 17, 17, 17, 3, 12, 4, 4, 14, 1, 0,
0, 0, 0, 0, 0, 5, 5, 5
])
self.assertEqual(ann_info['gt_labels_3d'].dtype, np.int64)
assert_allclose(ann_info['gt_labels_3d'], except_label)
self.assertIsInstance(ann_info['gt_bboxes_3d'], DepthInstance3DBoxes)
assert len(ann_info['gt_bboxes_3d']) == 27
assert torch.allclose(ann_info['gt_bboxes_3d'].tensor.sum(),
torch.tensor([107.7353]))
no_class_scannet_dataset = ScanNetDataset(
data_root, ann_file, metainfo=dict(CLASSES=['cabinet']))
input_dict = no_class_scannet_dataset.get_data_info(0)
ann_info = no_class_scannet_dataset.parse_ann_info(input_dict)
# assert the keys in ann_info and the type
self.assertIn('gt_labels_3d', ann_info)
# assert mapping to -1 or 1
assert (ann_info['gt_labels_3d'] <= 0).all()
self.assertEqual(ann_info['gt_labels_3d'].dtype, np.int64)
# all instance have been filtered by classes
self.assertEqual(len(ann_info['gt_labels_3d']), 27)
self.assertEqual(len(no_class_scannet_dataset.metainfo['CLASSES']), 1)
tests/test_data/test_datasets/test_sunrgbd_dataset.py 0 → 100644
View file @ 71b88ff5
# Copyright (c) OpenMMLab. All rights reserved.
import unittest
import numpy as np
import torch
from mmengine.testing import assert_allclose
from mmdet3d.core import DepthInstance3DBoxes
from mmdet3d.datasets import SUNRGBDDataset
def _generate_scannet_dataset_config():
data_root = 'tests/data/sunrgbd'
ann_file = 'sunrgbd_infos.pkl'
classes = ('bed', 'table', 'sofa', 'chair', 'toilet', 'desk', 'dresser',
'night_stand', 'bookshelf', 'bathtub')
from mmcv.transforms.base import BaseTransform
from mmengine.registry import TRANSFORMS
if 'Identity' not in TRANSFORMS:
@TRANSFORMS.register_module()
class Identity(BaseTransform):
def transform(self, info):
if 'ann_info' in info:
info['gt_labels_3d'] = info['ann_info']['gt_labels_3d']
return info
modality = dict(use_camera=True, use_lidar=True)
pipeline = [
dict(type='Identity'),
]
data_prefix = dict(pts='points', img='sunrgbd_trainval')
return data_root, ann_file, classes, data_prefix, pipeline, modality
class TestScanNetDataset(unittest.TestCase):
def test_sunrgbd_ataset(self):
np.random.seed(0)
data_root, ann_file, classes, data_prefix, \
pipeline, modality, = _generate_scannet_dataset_config()
scannet_dataset = SUNRGBDDataset(
data_root,
ann_file,
data_prefix=data_prefix,
pipeline=pipeline,
metainfo=dict(CLASSES=classes),
modality=modality)
scannet_dataset.prepare_data(0)
input_dict = scannet_dataset.get_data_info(0)
scannet_dataset[0]
# assert the the path should contains data_prefix and data_root
assert data_prefix['pts'] in input_dict['lidar_points']['lidar_path']
assert data_root in input_dict['lidar_points']['lidar_path']
for cam_id, img_info in input_dict['images'].items():
if 'img_path' in img_info:
assert data_prefix['img'] in img_info['img_path']
assert data_root in img_info['img_path']
ann_info = scannet_dataset.parse_ann_info(input_dict)
# assert the keys in ann_info and the type
except_label = np.array([0, 7, 6])
self.assertEqual(ann_info['gt_labels_3d'].dtype, np.int64)
assert_allclose(ann_info['gt_labels_3d'], except_label)
self.assertIsInstance(ann_info['gt_bboxes_3d'], DepthInstance3DBoxes)
self.assertEqual(len(ann_info['gt_bboxes_3d']), 3)
assert_allclose(ann_info['gt_bboxes_3d'].tensor.sum(),
torch.tensor(19.2575))
classes = ['bed']
bed_scannet_dataset = SUNRGBDDataset(
data_root,
ann_file,
data_prefix=data_prefix,
pipeline=pipeline,
metainfo=dict(CLASSES=classes),
modality=modality)
input_dict = bed_scannet_dataset.get_data_info(0)
ann_info = bed_scannet_dataset.parse_ann_info(input_dict)
# assert the keys in ann_info and the type
self.assertIn('gt_labels_3d', ann_info)
# assert mapping to -1 or 1
assert (ann_info['gt_labels_3d'] <= 0).all()
assert ann_info['gt_labels_3d'].dtype == np.int64
# all instance have been filtered by classes
self.assertEqual(len(ann_info['gt_labels_3d']), 3)
self.assertEqual(len(bed_scannet_dataset.metainfo['CLASSES']), 1)
tools/data_converter/update_infos_to_v2.py
View file @ 71b88ff5
......@@ -249,6 +249,15 @@ def update_kitti_infos(pkl_path, out_dir):
Trv2c = ori_info_dict['calib']['Tr_velo_to_cam'].astype(np.float32)
lidar2cam = rect @ Trv2c
temp_data_info['images']['CAM2']['lidar2cam'] = lidar2cam.tolist()
temp_data_info['images']['CAM0']['lidar2img'] = (
ori_info_dict['calib']['P0'] @ lidar2cam).tolist()
temp_data_info['images']['CAM1']['lidar2img'] = (
ori_info_dict['calib']['P1'] @ lidar2cam).tolist()
temp_data_info['images']['CAM2']['lidar2img'] = (
ori_info_dict['calib']['P2'] @ lidar2cam).tolist()
temp_data_info['images']['CAM3']['lidar2img'] = (
ori_info_dict['calib']['P3'] @ lidar2cam).tolist()
temp_data_info['lidar_points']['Tr_velo_to_cam'] = Trv2c.tolist()
# for potential usage
......@@ -311,6 +320,131 @@ def update_kitti_infos(pkl_path, out_dir):
mmcv.dump(converted_data_info, out_path, 'pkl')
def update_scannet_infos(pkl_path, out_dir):
print(f'{pkl_path} will be modified.')
if out_dir in pkl_path:
print(f'Warning, you may overwriting '
f'the original data {pkl_path}.')
time.sleep(5)
METAINFO = {
'CLASSES':
('cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window',
'bookshelf', 'picture', 'counter', 'desk', 'curtain', 'refrigerator',
'showercurtrain', 'toilet', 'sink', 'bathtub', 'garbagebin')
}
print(f'Reading from input file: {pkl_path}.')
data_list = mmcv.load(pkl_path)
print('Start updating:')
converted_list = []
for ori_info_dict in mmcv.track_iter_progress(data_list):
temp_data_info = get_empty_standard_data_info()
temp_data_info['lidar_points']['num_pts_feats'] = ori_info_dict[
'point_cloud']['num_features']
temp_data_info['lidar_points']['lidar_path'] = ori_info_dict[
'pts_path'].split('/')[-1]
temp_data_info['pts_semantic_mask_path'] = ori_info_dict[
'pts_semantic_mask_path'].split('/')[-1]
temp_data_info['pts_instance_mask_path'] = ori_info_dict[
'pts_instance_mask_path'].split('/')[-1]
# TODO support camera
# np.linalg.inv(info['axis_align_matrix'] @ extrinsic): depth2cam
anns = ori_info_dict['annos']
temp_data_info['axis_align_matrix'] = anns['axis_align_matrix'].tolist(
)
num_instances = len(anns['name'])
ignore_class_name = set()
instance_list = []
for instance_id in range(num_instances):
empty_instance = get_empty_instance()
empty_instance['bbox_3d'] = anns['gt_boxes_upright_depth'][
instance_id].tolist()
if anns['name'][instance_id] in METAINFO['CLASSES']:
empty_instance['bbox_label_3d'] = METAINFO['CLASSES'].index(
anns['name'][instance_id])
else:
ignore_class_name.add(anns['name'][instance_id])
empty_instance['bbox_label_3d'] = -1
empty_instance = clear_instance_unused_keys(empty_instance)
instance_list.append(empty_instance)
temp_data_info['instances'] = instance_list
temp_data_info, _ = clear_data_info_unused_keys(temp_data_info)
converted_list.append(temp_data_info)
pkl_name = pkl_path.split('/')[-1]
out_path = osp.join(out_dir, pkl_name)
print(f'Writing to output file: {out_path}.')
print(f'ignore classes: {ignore_class_name}')
converted_data_info = dict(
metainfo={'DATASET': 'SCANNET'}, data_list=converted_list)
mmcv.dump(converted_data_info, out_path, 'pkl')
def update_sunrgbd_infos(pkl_path, out_dir):
print(f'{pkl_path} will be modified.')
if out_dir in pkl_path:
print(f'Warning, you may overwriting '
f'the original data {pkl_path}.')
time.sleep(5)
METAINFO = {
'CLASSES': ('bed', 'table', 'sofa', 'chair', 'toilet', 'desk',
'dresser', 'night_stand', 'bookshelf', 'bathtub')
}
print(f'Reading from input file: {pkl_path}.')
data_list = mmcv.load(pkl_path)
print('Start updating:')
converted_list = []
for ori_info_dict in mmcv.track_iter_progress(data_list):
temp_data_info = get_empty_standard_data_info()
temp_data_info['lidar_points']['num_pts_feats'] = ori_info_dict[
'point_cloud']['num_features']
temp_data_info['lidar_points']['lidar_path'] = ori_info_dict[
'pts_path'].split('/')[-1]
calib = ori_info_dict['calib']
rt_mat = calib['Rt']
# follow Coord3DMode.convert_point
rt_mat = np.array([[1, 0, 0], [0, 0, -1], [0, 1, 0]
]) @ rt_mat.transpose(1, 0)
depth2img = calib['K'] @ rt_mat
temp_data_info['images']['CAM0']['depth2img'] = depth2img.tolist()
temp_data_info['images']['CAM0']['img_path'] = ori_info_dict['image'][
'image_path'].split('/')[-1]
h, w = ori_info_dict['image']['image_shape']
temp_data_info['images']['CAM0']['height'] = h
temp_data_info['images']['CAM0']['width'] = w
anns = ori_info_dict['annos']
num_instances = len(anns['name'])
ignore_class_name = set()
instance_list = []
for instance_id in range(num_instances):
empty_instance = get_empty_instance()
empty_instance['bbox_3d'] = anns['gt_boxes_upright_depth'][
instance_id].tolist()
if anns['name'][instance_id] in METAINFO['CLASSES']:
empty_instance['bbox_label_3d'] = METAINFO['CLASSES'].index(
anns['name'][instance_id])
else:
ignore_class_name.add(anns['name'][instance_id])
empty_instance['bbox_label_3d'] = -1
empty_instance = clear_instance_unused_keys(empty_instance)
instance_list.append(empty_instance)
temp_data_info['instances'] = instance_list
temp_data_info, _ = clear_data_info_unused_keys(temp_data_info)
converted_list.append(temp_data_info)
pkl_name = pkl_path.split('/')[-1]
out_path = osp.join(out_dir, pkl_name)
print(f'Writing to output file: {out_path}.')
print(f'ignore classes: {ignore_class_name}')
converted_data_info = dict(
metainfo={'DATASET': 'SUNRGBD'}, data_list=converted_list)
mmcv.dump(converted_data_info, out_path, 'pkl')
def parse_args():
parser = argparse.ArgumentParser(description='Arg parser for data coords '
'update due to coords sys refactor.')
......@@ -336,8 +470,15 @@ def main():
args = parse_args()
if args.out_dir is None:
args.out_dir = args.root_dir
if args.dataset == 'kitti':
if args.dataset.lower() == 'kitti':
update_kitti_infos(pkl_path=args.pkl, out_dir=args.out_dir)
elif args.dataset.lower() == 'scannet':
update_scannet_infos(pkl_path=args.pkl, out_dir=args.out_dir)
elif args.dataset.lower() == 'sunrgbd':
update_sunrgbd_infos(pkl_path=args.pkl, out_dir=args.out_dir)
else:
raise NotImplementedError(
f'Do not support convert {args.dataset} to v2.')
if __name__ == '__main__':
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment