Skip to content

GitLab

"...en/git@developer.sourcefind.cn:renzhc/diffusers_dcu.git" did not exist on "d816bcb5e836a2421f33e3f9ad10782993b9aecd"
Commit 360c27f9 authored by ZCMax's avatar ZCMax Committed by ChaimZhu
Browse files

[Refactor] Refactor 3D Seg Dataset

parent 1039ad0e
Hide whitespace changes
Inline Side-by-side
Showing with 1202 additions and 1261 deletions
configs/_base_/datasets/s3dis_seg-3d-13class.py
View file @ 360c27f9
# dataset settings
dataset_type = 'S3DISSegDataset'
data_root = './data/s3dis/'
# For S3DIS seg we usually do 13-class segmentation
class_names = ('ceiling', 'floor', 'wall', 'beam', 'column', 'window', 'door',
'table', 'chair', 'sofa', 'bookcase', 'board', 'clutter')
metainfo = dict(CLASSES=class_names)
dataset_type = 'S3DISSegDataset'
data_root = 'data/s3dis/'
input_modality = dict(use_lidar=True, use_camera=False)
data_prefix = dict(
pts='points',
pts_instance_mask='instance_mask',
pts_semantic_mask='semantic_mask')
file_client_args = dict(backend='disk')
# Uncomment the following if use ceph or other file clients.
# See https://mmcv.readthedocs.io/en/latest/api.html#mmcv.fileio.FileClient
......@@ -15,29 +22,27 @@ file_client_args = dict(backend='disk')
# 'data/s3dis/':
# 's3://openmmlab/datasets/detection3d/s3dis_processed/'
# }))
num_points = 4096
train_area = [1, 2, 3, 4, 6]
test_area = 5
train_pipeline = [
dict(
type='LoadPointsFromFile',
file_client_args=file_client_args,
coord_type='DEPTH',
shift_height=False,
use_color=True,
load_dim=6,
use_dim=[0, 1, 2, 3, 4, 5]),
use_dim=[0, 1, 2, 3, 4, 5],
file_client_args=file_client_args),
dict(
type='LoadAnnotations3D',
file_client_args=file_client_args,
with_bbox_3d=False,
with_label_3d=False,
with_mask_3d=False,
with_seg_3d=True),
dict(
type='PointSegClassMapping',
valid_cat_ids=tuple(range(len(class_names))),
max_cat_id=13),
with_seg_3d=True,
file_client_args=file_client_args),
dict(type='PointSegClassMapping'),
dict(
type='IndoorPatchPointSample',
num_points=num_points,
......@@ -47,18 +52,17 @@ train_pipeline = [
enlarge_size=0.2,
min_unique_num=None),
dict(type='NormalizePointsColor', color_mean=None),
dict(type='DefaultFormatBundle3D', class_names=class_names),
dict(type='Collect3D', keys=['points', 'pts_semantic_mask'])
dict(type='Pack3DDetInputs', keys=['points', 'pts_semantic_mask'])
]
test_pipeline = [
dict(
type='LoadPointsFromFile',
file_client_args=file_client_args,
coord_type='DEPTH',
shift_height=False,
use_color=True,
load_dim=6,
use_dim=[0, 1, 2, 3, 4, 5]),
use_dim=[0, 1, 2, 3, 4, 5],
file_client_args=file_client_args),
dict(type='NormalizePointsColor', color_mean=None),
dict(
# a wrapper in order to successfully call test function
......@@ -78,12 +82,8 @@ test_pipeline = [
sync_2d=False,
flip_ratio_bev_horizontal=0.0,
flip_ratio_bev_vertical=0.0),
dict(
type='DefaultFormatBundle3D',
class_names=class_names,
with_label=False),
dict(type='Collect3D', keys=['points'])
])
]),
dict(type='Pack3DDetInputs', keys=['points'])
]
# construct a pipeline for data and gt loading in show function
# please keep its loading function consistent with test_pipeline (e.g. client)
......@@ -91,69 +91,58 @@ test_pipeline = [
eval_pipeline = [
dict(
type='LoadPointsFromFile',
file_client_args=file_client_args,
coord_type='DEPTH',
shift_height=False,
use_color=True,
load_dim=6,
use_dim=[0, 1, 2, 3, 4, 5]),
dict(
type='LoadAnnotations3D',
file_client_args=file_client_args,
with_bbox_3d=False,
with_label_3d=False,
with_mask_3d=False,
with_seg_3d=True),
dict(
type='PointSegClassMapping',
valid_cat_ids=tuple(range(len(class_names))),
max_cat_id=13),
dict(
type='DefaultFormatBundle3D',
with_label=False,
class_names=class_names),
dict(type='Collect3D', keys=['points', 'pts_semantic_mask'])
use_dim=[0, 1, 2, 3, 4, 5],
file_client_args=file_client_args),
dict(type='NormalizePointsColor', color_mean=None),
dict(type='Pack3DDetInputs', keys=['points'])
]
data = dict(
samples_per_gpu=8,
workers_per_gpu=4,
# train on area 1, 2, 3, 4, 6
# test on area 5
train=dict(
# train on area 1, 2, 3, 4, 6
# test on area 5
train_dataloader = dict(
batch_size=8,
num_workers=4,
persistent_workers=True,
sampler=dict(type='DefaultSampler', shuffle=True),
dataset=dict(
type=dataset_type,
data_root=data_root,
ann_files=[
data_root + f's3dis_infos_Area_{i}.pkl' for i in train_area
],
metainfo=metainfo,
data_prefix=data_prefix,
pipeline=train_pipeline,
classes=class_names,
test_mode=False,
modality=input_modality,
ignore_index=len(class_names),
scene_idxs=[
data_root + f'seg_info/Area_{i}_resampled_scene_idxs.npy'
for i in train_area
],
file_client_args=file_client_args),
val=dict(
test_mode=False))
test_dataloader = dict(
batch_size=1,
num_workers=1,
persistent_workers=True,
drop_last=False,
sampler=dict(type='DefaultSampler', shuffle=False),
dataset=dict(
type=dataset_type,
data_root=data_root,
ann_files=data_root + f's3dis_infos_Area_{test_area}.pkl',
metainfo=metainfo,
data_prefix=data_prefix,
pipeline=test_pipeline,
classes=class_names,
test_mode=True,
modality=input_modality,
ignore_index=len(class_names),
scene_idxs=data_root +
f'seg_info/Area_{test_area}_resampled_scene_idxs.npy',
file_client_args=file_client_args),
test=dict(
type=dataset_type,
data_root=data_root,
ann_files=data_root + f's3dis_infos_Area_{test_area}.pkl',
pipeline=test_pipeline,
classes=class_names,
test_mode=True,
ignore_index=len(class_names),
file_client_args=file_client_args))
test_mode=True))
val_dataloader = test_dataloader
evaluation = dict(pipeline=eval_pipeline)
val_evaluator = dict(type='SegMetric')
test_evaluator = val_evaluator
configs/_base_/datasets/scannet_seg-3d-20class.py
View file @ 360c27f9
# dataset settings
dataset_type = 'ScanNetSegDataset'
data_root = './data/scannet/'
# For ScanNet seg we usually do 20-class segmentation
class_names = ('wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table',
'door', 'window', 'bookshelf', 'picture', 'counter', 'desk',
'curtain', 'refrigerator', 'showercurtrain', 'toilet', 'sink',
'bathtub', 'otherfurniture')
metainfo = dict(CLASSES=class_names)
dataset_type = 'ScanNetSegDataset'
data_root = 'data/scannet/'
input_modality = dict(use_lidar=True, use_camera=False)
data_prefix = dict(
pts='points',
pts_instance_mask='instance_mask',
pts_semantic_mask='semantic_mask')
file_client_args = dict(backend='disk')
# Uncomment the following if use ceph or other file clients.
# See https://mmcv.readthedocs.io/en/latest/api.html#mmcv.fileio.FileClient
# for more details.
# file_client_args = dict(
# backend='petrel',
# path_mapping=dict({
# './data/scannet/':
# 's3://openmmlab/datasets/detection3d/scannet_processed/',
# 'data/scannet/':
# 's3://openmmlab/datasets/detection3d/scannet_processed/'
# }))
num_points = 8192
train_pipeline = [
dict(
......@@ -13,18 +33,16 @@ train_pipeline = [
shift_height=False,
use_color=True,
load_dim=6,
use_dim=[0, 1, 2, 3, 4, 5]),
use_dim=[0, 1, 2, 3, 4, 5],
file_client_args=file_client_args),
dict(
type='LoadAnnotations3D',
with_bbox_3d=False,
with_label_3d=False,
with_mask_3d=False,
with_seg_3d=True),
dict(
type='PointSegClassMapping',
valid_cat_ids=(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28,
33, 34, 36, 39),
max_cat_id=40),
with_seg_3d=True,
file_client_args=file_client_args),
dict(type='PointSegClassMapping'),
dict(
type='IndoorPatchPointSample',
num_points=num_points,
......@@ -34,8 +52,7 @@ train_pipeline = [
enlarge_size=0.2,
min_unique_num=None),
dict(type='NormalizePointsColor', color_mean=None),
dict(type='DefaultFormatBundle3D', class_names=class_names),
dict(type='Collect3D', keys=['points', 'pts_semantic_mask'])
dict(type='Pack3DDetInputs', keys=['points', 'pts_semantic_mask'])
]
test_pipeline = [
dict(
......@@ -44,7 +61,8 @@ test_pipeline = [
shift_height=False,
use_color=True,
load_dim=6,
use_dim=[0, 1, 2, 3, 4, 5]),
use_dim=[0, 1, 2, 3, 4, 5],
file_client_args=file_client_args),
dict(type='NormalizePointsColor', color_mean=None),
dict(
# a wrapper in order to successfully call test function
......@@ -64,12 +82,8 @@ test_pipeline = [
sync_2d=False,
flip_ratio_bev_horizontal=0.0,
flip_ratio_bev_vertical=0.0),
dict(
type='DefaultFormatBundle3D',
class_names=class_names,
with_label=False),
dict(type='Collect3D', keys=['points'])
])
]),
dict(type='Pack3DDetInputs', keys=['points'])
]
# construct a pipeline for data and gt loading in show function
# please keep its loading function consistent with test_pipeline (e.g. client)
......@@ -81,52 +95,45 @@ eval_pipeline = [
shift_height=False,
use_color=True,
load_dim=6,
use_dim=[0, 1, 2, 3, 4, 5]),
dict(
type='LoadAnnotations3D',
with_bbox_3d=False,
with_label_3d=False,
with_mask_3d=False,
with_seg_3d=True),
dict(
type='PointSegClassMapping',
valid_cat_ids=(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28,
33, 34, 36, 39),
max_cat_id=40),
dict(
type='DefaultFormatBundle3D',
with_label=False,
class_names=class_names),
dict(type='Collect3D', keys=['points', 'pts_semantic_mask'])
use_dim=[0, 1, 2, 3, 4, 5],
file_client_args=file_client_args),
dict(type='NormalizePointsColor', color_mean=None),
dict(type='Pack3DDetInputs', keys=['points'])
]
data = dict(
samples_per_gpu=8,
workers_per_gpu=4,
train=dict(
train_dataloader = dict(
batch_size=8,
num_workers=4,
persistent_workers=True,
sampler=dict(type='DefaultSampler', shuffle=True),
dataset=dict(
type=dataset_type,
data_root=data_root,
ann_file=data_root + 'scannet_infos_train.pkl',
ann_file='scannet_infos_train.pkl',
metainfo=metainfo,
data_prefix=data_prefix,
pipeline=train_pipeline,
classes=class_names,
test_mode=False,
modality=input_modality,
ignore_index=len(class_names),
scene_idxs=data_root + 'seg_info/train_resampled_scene_idxs.npy'),
val=dict(
scene_idxs=data_root + 'seg_info/train_resampled_scene_idxs.npy',
test_mode=False))
test_dataloader = dict(
batch_size=1,
num_workers=1,
persistent_workers=True,
drop_last=False,
sampler=dict(type='DefaultSampler', shuffle=False),
dataset=dict(
type=dataset_type,
data_root=data_root,
ann_file=data_root + 'scannet_infos_val.pkl',
ann_file='scannet_infos_val.pkl',
metainfo=metainfo,
data_prefix=data_prefix,
pipeline=test_pipeline,
classes=class_names,
test_mode=True,
ignore_index=len(class_names)),
test=dict(
type=dataset_type,
data_root=data_root,
ann_file=data_root + 'scannet_infos_val.pkl',
pipeline=test_pipeline,
classes=class_names,
test_mode=True,
ignore_index=len(class_names)))
modality=input_modality,
ignore_index=len(class_names),
test_mode=True))
val_dataloader = test_dataloader
evaluation = dict(pipeline=eval_pipeline)
val_evaluator = dict(type='SegMetric')
test_evaluator = val_evaluator
mmdet3d/core/data_structures/__init__.py
View file @ 360c27f9
# Copyright (c) OpenMMLab. All rights reserved.
from .det3d_data_sample import Det3DDataSample
from .point_data import PointData
__all__ = ['Det3DDataSample']
__all__ = ['Det3DDataSample', 'PointData']
mmdet3d/core/data_structures/det3d_data_sample.py
View file @ 360c27f9
# Copyright (c) OpenMMLab. All rights reserved.
from mmengine.data import InstanceData, PixelData
from mmengine.data import InstanceData
from mmdet.core.data_structures import DetDataSample
from .point_data import PointData
class Det3DDataSample(DetDataSample):
......@@ -43,19 +44,15 @@ class Det3DDataSample(DetDataSample):
`use_lidar=True, use_camera=True`, the 3D predictions based on
image are saved in `img_pred_instances_3d` to distinguish with
`pts_pred_instances_3d` which based on point cloud.
- ``gt_pts_sem_seg``(PixelData): Ground truth of point cloud
semantic segmentation.
- ``pred_pts_sem_seg``(PixelData): Prediction of point cloud
semantic segmentation.
- ``gt_pts_panoptic_seg``(PixelData): Ground truth of point cloud
panoptic segmentation.
- ``pred_pts_panoptic_seg``(PixelData): Predicted of point cloud
panoptic segmentation.
- ``gt_pts_seg``(PointData): Ground truth of point cloud
segmentation.
- ``pred_pts_seg``(PointData): Prediction of point cloud
segmentation.
- ``eval_ann_info``(dict): Raw annotation, which will be passed to
evaluator and do the online evaluation.
Examples:
>>> from mmengine.data import InstanceData, PixelData
>>> from mmengine.data import InstanceData
>>> from mmdet3d.core import Det3DDataSample
>>> from mmdet3d.core.bbox import BaseInstance3DBoxes
......@@ -128,38 +125,33 @@ class Det3DDataSample(DetDataSample):
>>> assert 'bboxes' in data_sample.gt_instances_3d
>>> data_sample = Det3DDataSample()
>>> gt_pts_panoptic_seg_data = dict(panoptic_seg=torch.rand(1, 2, 4))
>>> gt_pts_panoptic_seg = PixelData(**gt_pts_panoptic_seg_data)
>>> data_sample.gt_pts_panoptic_seg = gt_pts_panoptic_seg
... gt_pts_seg_data = dict(
... pts_instance_mask=torch.rand(2),
... pts_semantic_mask=torch.rand(2))
>>> data_sample.gt_pts_seg = PointData(**gt_pts_seg_data)
>>> print(data_sample)
<Det3DDataSample(
META INFORMATION
DATA FIELDS
_gt_pts_panoptic_seg: <PixelData(
gt_pts_seg: <PointData(
META INFORMATION
DATA FIELDS
panoptic_seg: tensor([[[0.9875, 0.3012, 0.5534, 0.9593],
[0.1251, 0.1911, 0.8058, 0.2566]]])
) at 0x7fb0d93543d0>
gt_pts_panoptic_seg: <PixelData(
pts_instance_mask: tensor([0.0576, 0.3067])
pts_semantic_mask: tensor([0.9267, 0.7455])
) at 0x7f654a9c1590>
_gt_pts_seg: <PointData(
META INFORMATION
DATA FIELDS
panoptic_seg: tensor([[[0.9875, 0.3012, 0.5534, 0.9593],
[0.1251, 0.1911, 0.8058, 0.2566]]])
) at 0x7fb0d93543d0>
) at 0x7fb0d9354280>
>>> data_sample = Det3DDataSample()
>>> gt_pts_sem_seg_data = dict(segm_seg=torch.rand(2, 2, 2))
>>> gt_pts_sem_seg = PixelData(**gt_pts_sem_seg_data)
>>> data_sample.gt_pts_sem_seg = gt_pts_sem_seg
>>> assert 'gt_pts_sem_seg' in data_sample
>>> assert 'segm_seg' in data_sample.gt_pts_sem_seg
pts_instance_mask: tensor([0.0576, 0.3067])
pts_semantic_mask: tensor([0.9267, 0.7455])
) at 0x7f654a9c1590>
) at 0x7f654a9c1550>
"""
@property
......@@ -211,49 +203,25 @@ class Det3DDataSample(DetDataSample):
del self._img_pred_instances_3d
@property
def gt_pts_sem_seg(self) -> PixelData:
return self._gt_pts_sem_seg
@gt_pts_sem_seg.setter
def gt_pts_sem_seg(self, value: PixelData):
self.set_field(value, '_gt_pts_sem_seg', dtype=PixelData)
@gt_pts_sem_seg.deleter
def gt_pts_sem_seg(self):
del self._gt_pts_sem_seg
@property
def pred_pts_sem_seg(self) -> PixelData:
return self._pred_pts_sem_seg
@pred_pts_sem_seg.setter
def pred_pts_sem_seg(self, value: PixelData):
self.set_field(value, '_pred_pts_sem_seg', dtype=PixelData)
@pred_pts_sem_seg.deleter
def pred_pts_sem_seg(self):
del self._pred_pts_sem_seg
@property
def gt_pts_panoptic_seg(self) -> PixelData:
return self._gt_pts_panoptic_seg
def gt_pts_seg(self) -> PointData:
return self._gt_pts_seg
@gt_pts_panoptic_seg.setter
def gt_pts_panoptic_seg(self, value: PixelData):
self.set_field(value, '_gt_pts_panoptic_seg', dtype=PixelData)
@gt_pts_seg.setter
def gt_pts_seg(self, value: PointData):
self.set_field(value, '_gt_pts_seg', dtype=PointData)
@gt_pts_panoptic_seg.deleter
def gt_pts_panoptic_seg(self):
del self._gt_pts_panoptic_seg
@gt_pts_seg.deleter
def gt_pts_seg(self):
del self._gt_pts_seg
@property
def pred_pts_panoptic_seg(self) -> PixelData:
return self._pred_pts_panoptic_seg
def pred_pts_seg(self) -> PointData:
return self._pred_pts_seg
@pred_pts_panoptic_seg.setter
def pred_pts_panoptic_seg(self, value: PixelData):
self.set_field(value, '_pred_pts_panoptic_seg', dtype=PixelData)
@pred_pts_seg.setter
def pred_pts_seg(self, value: PointData):
self.set_field(value, '_pred_pts_seg', dtype=PointData)
@pred_pts_panoptic_seg.deleter
def pred_pts_panoptic_seg(self):
del self._pred_pts_panoptic_seg
@pred_pts_seg.deleter
def pred_pts_seg(self):
del self._pred_pts_seg
mmdet3d/core/data_structures/point_data.py 0 → 100644
View file @ 360c27f9
# Copyright (c) OpenMMLab. All rights reserved.
from collections.abc import Sized
from typing import Union
import numpy as np
import torch
from mmengine.data import BaseDataElement
IndexType = Union[str, slice, int, list, torch.LongTensor,
torch.cuda.LongTensor, torch.BoolTensor,
torch.cuda.BoolTensor, np.ndarray]
class PointData(BaseDataElement):
"""Data structure for point-level annnotations or predictions.
All data items in ``data_fields`` of ``PointData`` meet the following
requirements:
- They are all one dimension.
- They should have the same length.
Notice: ``PointData`` behaves like `InstanceData`.
Examples:
>>> metainfo = dict(
... sample_id=random.randint(0, 100))
>>> points = np.random.randint(0, 255, (100, 3))
>>> point_data = PointData(metainfo=metainfo,
... points=points)
>>> print(len(point_data))
>>> (100)
>>> # slice
>>> slice_data = pixel_data[10:60]
>>> assert slice_data.shape == (50,)
>>> # set
>>> point_data.pts_semantic_mask = torch.randint(0, 255, (100))
>>> point_data.pts_instance_mask = torch.randint(0, 255, (100))
>>> assert tuple(point_data.pts_semantic_mask.shape) == (100)
>>> assert tuple(point_data.pts_instance_mask.shape) == (100)
"""
def __setattr__(self, name: str, value: Sized):
"""setattr is only used to set data.
the value must have the attribute of `__len__` and have the same length
of PointData.
"""
if name in ('_metainfo_fields', '_data_fields'):
if not hasattr(self, name):
super().__setattr__(name, value)
else:
raise AttributeError(
f'{name} has been used as a '
f'private attribute, which is immutable. ')
else:
assert isinstance(value,
Sized), 'value must contain `_len__` attribute'
if len(self) > 0:
assert len(value) == len(self), f'the length of ' \
f'values {len(value)} is ' \
f'not consistent with' \
f' the length of this ' \
f':obj:`PointData` ' \
f'{len(self)} '
super().__setattr__(name, value)
__setitem__ = __setattr__
def __getitem__(self, item: IndexType) -> 'PointData':
"""
Args:
item (str, obj:`slice`,
obj`torch.LongTensor`, obj:`torch.BoolTensor`):
get the corresponding values according to item.
Returns:
obj:`PointData`: Corresponding values.
"""
if isinstance(item, list):
item = np.array(item)
if isinstance(item, np.ndarray):
item = torch.from_numpy(item)
assert isinstance(
item, (str, slice, int, torch.LongTensor, torch.cuda.LongTensor,
torch.BoolTensor, torch.cuda.BoolTensor))
if isinstance(item, str):
return getattr(self, item)
if type(item) == int:
if item >= len(self) or item < -len(self): # type:ignore
raise IndexError(f'Index {item} out of range!')
else:
# keep the dimension
item = slice(item, None, len(self))
new_data = self.__class__(metainfo=self.metainfo)
if isinstance(item, torch.Tensor):
assert item.dim() == 1, 'Only support to get the' \
' values along the first dimension.'
if isinstance(item, (torch.BoolTensor, torch.cuda.BoolTensor)):
assert len(item) == len(self), f'The shape of the' \
f' input(BoolTensor)) ' \
f'{len(item)} ' \
f' does not match the shape ' \
f'of the indexed tensor ' \
f'in results_filed ' \
f'{len(self)} at ' \
f'first dimension. '
for k, v in self.items():
if isinstance(v, torch.Tensor):
new_data[k] = v[item]
elif isinstance(v, np.ndarray):
new_data[k] = v[item.cpu().numpy()]
elif isinstance(
v, (str, list, tuple)) or (hasattr(v, '__getitem__')
and hasattr(v, 'cat')):
# convert to indexes from boolTensor
if isinstance(item,
(torch.BoolTensor, torch.cuda.BoolTensor)):
indexes = torch.nonzero(item).view(
-1).cpu().numpy().tolist()
else:
indexes = item.cpu().numpy().tolist()
slice_list = []
if indexes:
for index in indexes:
slice_list.append(slice(index, None, len(v)))
else:
slice_list.append(slice(None, 0, None))
r_list = [v[s] for s in slice_list]
if isinstance(v, (str, list, tuple)):
new_value = r_list[0]
for r in r_list[1:]:
new_value = new_value + r
else:
new_value = v.cat(r_list)
new_data[k] = new_value
else:
raise ValueError(
f'The type of `{k}` is `{type(v)}`, which has no '
'attribute of `cat`, so it does not '
f'support slice with `bool`')
else:
# item is a slice
for k, v in self.items():
new_data[k] = v[item]
return new_data # type:ignore
def __len__(self) -> int:
"""int: the length of PointData"""
if len(self._data_fields) > 0:
return len(self.values()[0])
else:
return 0
mmdet3d/datasets/__init__.py
View file @ 360c27f9
# Copyright (c) OpenMMLab. All rights reserved.
from .builder import DATASETS, PIPELINES, build_dataset
from .custom_3d_seg import Custom3DSegDataset
from .det3d_dataset import Det3DDataset
from .kitti_dataset import KittiDataset
from .kitti_mono_dataset import KittiMonoDataset
......@@ -22,6 +21,7 @@ from .pipelines import (AffineResize, BackgroundPointsFilter, GlobalAlignment,
from .s3dis_dataset import S3DISDataset, S3DISSegDataset
from .scannet_dataset import (ScanNetDataset, ScanNetInstanceSegDataset,
ScanNetSegDataset)
from .seg3d_dataset import Seg3DDataset
from .semantickitti_dataset import SemanticKITTIDataset
from .sunrgbd_dataset import SUNRGBDDataset
from .utils import get_loading_pipeline
......@@ -36,7 +36,7 @@ __all__ = [
'IndoorPatchPointSample', 'IndoorPointSample', 'PointSample',
'LoadAnnotations3D', 'GlobalAlignment', 'SUNRGBDDataset', 'ScanNetDataset',
'ScanNetSegDataset', 'ScanNetInstanceSegDataset', 'SemanticKITTIDataset',
'Det3DDataset', 'Custom3DSegDataset', 'LoadPointsFromMultiSweeps',
'Det3DDataset', 'Seg3DDataset', 'LoadPointsFromMultiSweeps',
'WaymoDataset', 'BackgroundPointsFilter', 'VoxelBasedPointSampler',
'get_loading_pipeline', 'RandomDropPointsColor', 'RandomJitterPoints',
'ObjectNameFilter', 'AffineResize', 'RandomShiftScale',
......
mmdet3d/datasets/custom_3d_seg.py deleted 100644 → 0
View file @ 1039ad0e
# Copyright (c) OpenMMLab. All rights reserved.
import tempfile
import warnings
from os import path as osp
import mmcv
import numpy as np
from torch.utils.data import Dataset
from mmdet3d.registry import DATASETS
from .pipelines import Compose
from .utils import extract_result_dict, get_loading_pipeline
@DATASETS.register_module()
class Custom3DSegDataset(Dataset):
"""Customized 3D dataset for semantic segmentation task.
This is the base dataset of ScanNet and S3DIS dataset.
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.
Defaults to None.
palette (list[list[int]], optional): The palette of segmentation map.
Defaults to None.
modality (dict, optional): Modality to specify the sensor data used
as input. Defaults to None.
test_mode (bool, optional): Whether the dataset is in test mode.
Defaults to False.
ignore_index (int, optional): The label index to be ignored, e.g.
unannotated points. If None is given, set to len(self.CLASSES) to
be consistent with PointSegClassMapping function in pipeline.
Defaults to None.
scene_idxs (np.ndarray | str, optional): Precomputed index to load
data. For scenes with many points, we may sample it several times.
Defaults to None.
"""
# names of all classes data used for the task
CLASSES = None
# class_ids used for training
VALID_CLASS_IDS = None
# all possible class_ids in loaded segmentation mask
ALL_CLASS_IDS = None
# official color for visualization
PALETTE = None
def __init__(self,
data_root,
ann_file,
pipeline=None,
classes=None,
palette=None,
modality=None,
test_mode=False,
ignore_index=None,
scene_idxs=None,
file_client_args=dict(backend='disk')):
super().__init__()
self.data_root = data_root
self.ann_file = ann_file
self.test_mode = test_mode
self.modality = modality
self.file_client = mmcv.FileClient(**file_client_args)
# load annotations
if hasattr(self.file_client, 'get_local_path'):
with self.file_client.get_local_path(self.ann_file) as local_path:
self.data_infos = self.load_annotations(open(local_path, 'rb'))
else:
warnings.warn(
'The used MMCV version does not have get_local_path. '
f'We treat the {self.ann_file} as local paths and it '
'might cause errors if the path is not a local path. '
'Please use MMCV>= 1.3.16 if you meet errors.')
self.data_infos = self.load_annotations(self.ann_file)
if pipeline is not None:
self.pipeline = Compose(pipeline)
self.ignore_index = len(self.CLASSES) if \
ignore_index is None else ignore_index
self.scene_idxs = self.get_scene_idxs(scene_idxs)
self.CLASSES, self.PALETTE = \
self.get_classes_and_palette(classes, palette)
# set group flag for the sampler
if not self.test_mode:
self._set_group_flag()
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.
"""
# loading data from a file-like object needs file format
return mmcv.load(ann_file, file_format='pkl')
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.
- file_name (str): Filename of point clouds.
- 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(
pts_filename=pts_filename,
sample_idx=sample_idx,
file_name=pts_filename)
if not self.test_mode:
annos = self.get_ann_info(index)
input_dict['ann_info'] = annos
return input_dict
def pre_pipeline(self, results):
"""Initialization before data preparation.
Args:
results (dict): Dict before data preprocessing.
- img_fields (list): Image fields.
- pts_mask_fields (list): Mask fields of points.
- pts_seg_fields (list): Mask fields of point segments.
- mask_fields (list): Fields of masks.
- seg_fields (list): Segment fields.
"""
results['img_fields'] = []
results['pts_mask_fields'] = []
results['pts_seg_fields'] = []
results['mask_fields'] = []
results['seg_fields'] = []
results['bbox3d_fields'] = []
def prepare_train_data(self, index):
"""Training data preparation.
Args:
index (int): Index for accessing the target data.
Returns:
dict: Training data dict of the corresponding index.
"""
input_dict = self.get_data_info(index)
if input_dict is None:
return None
self.pre_pipeline(input_dict)
example = self.pipeline(input_dict)
return example
def prepare_test_data(self, index):
"""Prepare data for testing.
Args:
index (int): Index for accessing the target data.
Returns:
dict: Testing data dict of the corresponding index.
"""
input_dict = self.get_data_info(index)
self.pre_pipeline(input_dict)
example = self.pipeline(input_dict)
return example
def get_classes_and_palette(self, classes=None, palette=None):
"""Get class names of current dataset.
This function is taken from MMSegmentation.
Args:
classes (Sequence[str] | str): If classes is None, use
default CLASSES defined by builtin dataset. If classes is a
string, take it as a file name. The file contains the name of
classes where each line contains one class name. If classes is
a tuple or list, override the CLASSES defined by the dataset.
Defaults to None.
palette (Sequence[Sequence[int]]] | np.ndarray):
The palette of segmentation map. If None is given, random
palette will be generated. Defaults to None.
"""
if classes is None:
self.custom_classes = False
# map id in the loaded mask to label used for training
self.label_map = {
cls_id: self.ignore_index
for cls_id in self.ALL_CLASS_IDS
}
self.label_map.update(
{cls_id: i
for i, cls_id in enumerate(self.VALID_CLASS_IDS)})
# map label to category name
self.label2cat = {
i: cat_name
for i, cat_name in enumerate(self.CLASSES)
}
return self.CLASSES, self.PALETTE
self.custom_classes = True
if isinstance(classes, str):
# take it as a file path
class_names = mmcv.list_from_file(classes)
elif isinstance(classes, (tuple, list)):
class_names = classes
else:
raise ValueError(f'Unsupported type {type(classes)} of classes.')
if self.CLASSES:
if not set(class_names).issubset(self.CLASSES):
raise ValueError('classes is not a subset of CLASSES.')
# update valid_class_ids
self.VALID_CLASS_IDS = [
self.VALID_CLASS_IDS[self.CLASSES.index(cls_name)]
for cls_name in class_names
]
# dictionary, its keys are the old label ids and its values
# are the new label ids.
# used for changing pixel labels in load_annotations.
self.label_map = {
cls_id: self.ignore_index
for cls_id in self.ALL_CLASS_IDS
}
self.label_map.update(
{cls_id: i
for i, cls_id in enumerate(self.VALID_CLASS_IDS)})
self.label2cat = {
i: cat_name
for i, cat_name in enumerate(class_names)
}
# modify palette for visualization
palette = [
self.PALETTE[self.CLASSES.index(cls_name)]
for cls_name in class_names
]
return class_names, palette
def get_scene_idxs(self, scene_idxs):
"""Compute scene_idxs for data sampling.
We sample more times for scenes with more points.
"""
if self.test_mode:
# when testing, we load one whole scene every time
return np.arange(len(self.data_infos)).astype(np.int32)
# we may need to re-sample different scenes according to scene_idxs
# this is necessary for indoor scene segmentation such as ScanNet
if scene_idxs is None:
scene_idxs = np.arange(len(self.data_infos))
if isinstance(scene_idxs, str):
with self.file_client.get_local_path(scene_idxs) as local_path:
scene_idxs = np.load(local_path)
else:
scene_idxs = np.array(scene_idxs)
return scene_idxs.astype(np.int32)
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): 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.
Returns:
tuple: (outputs, tmp_dir), outputs is the detection results,
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')
out = f'{pklfile_prefix}.pkl'
mmcv.dump(outputs, out)
return outputs, tmp_dir
def evaluate(self,
results,
metric=None,
logger=None,
show=False,
out_dir=None,
pipeline=None):
"""Evaluate.
Evaluation in semantic segmentation protocol.
Args:
results (list[dict]): List of results.
metric (str | list[str]): Metrics to be evaluated.
logger (logging.Logger | str, optional): Logger used for printing
related information during evaluation. Defaults to None.
show (bool, optional): Whether to visualize.
Defaults to False.
out_dir (str, optional): Path to save the visualization results.
Defaults to None.
pipeline (list[dict], optional): raw data loading for showing.
Default: None.
Returns:
dict: Evaluation results.
"""
from mmdet3d.core.evaluation import seg_eval
assert isinstance(
results, list), f'Expect results to be list, got {type(results)}.'
assert len(results) > 0, 'Expect length of results > 0.'
assert len(results) == len(self.data_infos)
assert isinstance(
results[0], dict
), f'Expect elements in results to be dict, got {type(results[0])}.'
load_pipeline = self._get_pipeline(pipeline)
pred_sem_masks = [result['semantic_mask'] for result in results]
gt_sem_masks = [
self._extract_data(
i, load_pipeline, 'pts_semantic_mask', load_annos=True)
for i in range(len(self.data_infos))
]
ret_dict = seg_eval(
gt_sem_masks,
pred_sem_masks,
self.label2cat,
self.ignore_index,
logger=logger)
if show:
self.show(pred_sem_masks, out_dir, pipeline=pipeline)
return ret_dict
def _rand_another(self, idx):
"""Randomly get another item with the same flag.
Returns:
int: Another index of item with the same flag.
"""
pool = np.where(self.flag == self.flag[idx])[0]
return np.random.choice(pool)
def _build_default_pipeline(self):
"""Build the default pipeline for this dataset."""
raise NotImplementedError('_build_default_pipeline is not implemented '
f'for dataset {self.__class__.__name__}')
def _get_pipeline(self, pipeline):
"""Get data loading pipeline in self.show/evaluate function.
Args:
pipeline (list[dict]): Input pipeline. If None is given,
get from self.pipeline.
"""
if pipeline is None:
if not hasattr(self, 'pipeline') or self.pipeline is None:
warnings.warn(
'Use default pipeline for data loading, this may cause '
'errors when data is on ceph')
return self._build_default_pipeline()
loading_pipeline = get_loading_pipeline(self.pipeline.transforms)
return Compose(loading_pipeline)
return Compose(pipeline)
def _extract_data(self, index, pipeline, key, load_annos=False):
"""Load data using input pipeline and extract data according to key.
Args:
index (int): Index for accessing the target data.
pipeline (:obj:`Compose`): Composed data loading pipeline.
key (str | list[str]): One single or a list of data key.
load_annos (bool): Whether to load data annotations.
If True, need to set self.test_mode as False before loading.
Returns:
np.ndarray | torch.Tensor | list[np.ndarray | torch.Tensor]:
A single or a list of loaded data.
"""
assert pipeline is not None, 'data loading pipeline is not provided'
# when we want to load ground-truth via pipeline (e.g. bbox, seg mask)
# we need to set self.test_mode as False so that we have 'annos'
if load_annos:
original_test_mode = self.test_mode
self.test_mode = False
input_dict = self.get_data_info(index)
self.pre_pipeline(input_dict)
example = pipeline(input_dict)
# extract data items according to keys
if isinstance(key, str):
data = extract_result_dict(example, key)
else:
data = [extract_result_dict(example, k) for k in key]
if load_annos:
self.test_mode = original_test_mode
return data
def __len__(self):
"""Return the length of scene_idxs.
Returns:
int: Length of data infos.
"""
return len(self.scene_idxs)
def __getitem__(self, idx):
"""Get item from infos according to the given index.
In indoor scene segmentation task, each scene contains millions of
points. However, we only sample less than 10k points within a patch
each time. Therefore, we use `scene_idxs` to re-sample different rooms.
Returns:
dict: Data dictionary of the corresponding index.
"""
scene_idx = self.scene_idxs[idx] # map to scene idx
if self.test_mode:
return self.prepare_test_data(scene_idx)
while True:
data = self.prepare_train_data(scene_idx)
if data is None:
idx = self._rand_another(idx)
scene_idx = self.scene_idxs[idx] # map to scene idx
continue
return data
def _set_group_flag(self):
"""Set flag according to image aspect ratio.
Images with aspect ratio greater than 1 will be set as group 1,
otherwise group 0. In 3D datasets, they are all the same, thus are all
zeros.
"""
self.flag = np.zeros(len(self), dtype=np.uint8)
mmdet3d/datasets/pipelines/formating.py
View file @ 360c27f9
......@@ -6,7 +6,7 @@ from mmcv import BaseTransform
from mmcv.transforms import to_tensor
from mmengine import InstanceData
from mmdet3d.core import Det3DDataSample
from mmdet3d.core import Det3DDataSample, PointData
from mmdet3d.core.bbox import BaseInstance3DBoxes
from mmdet3d.core.points import BasePoints
from mmdet3d.registry import TRANSFORMS
......@@ -143,7 +143,7 @@ class Pack3DDetInputs(BaseTransform):
data_sample = Det3DDataSample()
gt_instances_3d = InstanceData()
gt_instances = InstanceData()
seg_data = dict()
gt_pts_seg = PointData()
img_metas = {}
for key in self.meta_keys:
......@@ -161,7 +161,7 @@ class Pack3DDetInputs(BaseTransform):
elif key in self.INSTANCEDATA_2D_KEYS:
gt_instances[self._remove_prefix(key)] = results[key]
elif key in self.SEG_KEYS:
seg_data[self._remove_prefix(key)] = results[key]
gt_pts_seg[self._remove_prefix(key)] = results[key]
else:
raise NotImplementedError(f'Please modified '
f'`Pack3DDetInputs` '
......@@ -170,7 +170,7 @@ class Pack3DDetInputs(BaseTransform):
data_sample.gt_instances_3d = gt_instances_3d
data_sample.gt_instances = gt_instances
data_sample.seg_data = seg_data
data_sample.gt_pts_seg = gt_pts_seg
if 'eval_ann_info' in results:
data_sample.eval_ann_info = results['eval_ann_info']
else:
......
mmdet3d/datasets/pipelines/loading.py
View file @ 360c27f9
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Sequence
from typing import List
import mmcv
import numpy as np
......@@ -270,22 +270,6 @@ class PointSegClassMapping(BaseTransform):
segmentation mask. Defaults to 40.
"""
def __init__(self,
valid_cat_ids: Sequence[int],
max_cat_id: int = 40) -> None:
assert max_cat_id >= np.max(valid_cat_ids), \
'max_cat_id should be greater than maximum id in valid_cat_ids'
self.valid_cat_ids = valid_cat_ids
self.max_cat_id = int(max_cat_id)
# build cat_id to class index mapping
neg_cls = len(valid_cat_ids)
self.cat_id2class = np.ones(
self.max_cat_id + 1, dtype=np.int) * neg_cls
for cls_idx, cat_id in enumerate(valid_cat_ids):
self.cat_id2class[cat_id] = cls_idx
def transform(self, results: dict) -> None:
"""Call function to map original semantic class to valid category ids.
......@@ -301,9 +285,19 @@ class PointSegClassMapping(BaseTransform):
assert 'pts_semantic_mask' in results
pts_semantic_mask = results['pts_semantic_mask']
converted_pts_sem_mask = self.cat_id2class[pts_semantic_mask]
assert 'label_mapping' in results
label_mapping = results['label_mapping']
converted_pts_sem_mask = \
np.array([label_mapping[mask] for mask in pts_semantic_mask])
results['pts_semantic_mask'] = converted_pts_sem_mask
# 'eval_ann_info' will be passed to evaluator
if 'eval_ann_info' in results:
assert 'pts_semantic_mask' in results['eval_ann_info']
results['eval_ann_info']['pts_semantic_mask'] = \
converted_pts_sem_mask
return results
def __repr__(self):
......@@ -315,17 +309,17 @@ class PointSegClassMapping(BaseTransform):
@TRANSFORMS.register_module()
class NormalizePointsColor(object):
class NormalizePointsColor(BaseTransform):
"""Normalize color of points.
Args:
color_mean (list[float]): Mean color of the point cloud.
"""
def __init__(self, color_mean):
def __init__(self, color_mean: List[float]) -> None:
self.color_mean = color_mean
def __call__(self, results):
def transform(self, input_dict: dict) -> dict:
"""Call function to normalize color of points.
Args:
......@@ -337,7 +331,7 @@ class NormalizePointsColor(object):
- points (:obj:`BasePoints`): Points after color normalization.
"""
points = results['points']
points = input_dict['points']
assert points.attribute_dims is not None and \
'color' in points.attribute_dims.keys(), \
'Expect points have color attribute'
......@@ -345,8 +339,8 @@ class NormalizePointsColor(object):
points.color = points.color - \
points.color.new_tensor(self.color_mean)
points.color = points.color / 255.0
results['points'] = points
return results
input_dict['points'] = points
return input_dict
def __repr__(self):
"""str: Return a string that describes the module."""
......
mmdet3d/datasets/pipelines/transforms_3d.py
View file @ 360c27f9
# Copyright (c) OpenMMLab. All rights reserved.
import random
import warnings
from typing import Dict, List
from typing import Dict, List, Optional, Tuple, Union
import cv2
import numpy as np
......@@ -18,7 +18,7 @@ from .data_augment_utils import noise_per_object_v3_
@TRANSFORMS.register_module()
class RandomDropPointsColor(object):
class RandomDropPointsColor(BaseTransform):
r"""Randomly set the color of points to all zeros.
Once this transform is executed, all the points' color will be dropped.
......@@ -30,12 +30,12 @@ class RandomDropPointsColor(object):
Defaults to 0.2.
"""
def __init__(self, drop_ratio=0.2):
def __init__(self, drop_ratio: float = 0.2) -> None:
assert isinstance(drop_ratio, (int, float)) and 0 <= drop_ratio <= 1, \
f'invalid drop_ratio value {drop_ratio}'
self.drop_ratio = drop_ratio
def __call__(self, input_dict):
def transform(self, input_dict: dict) -> dict:
"""Call function to drop point colors.
Args:
......@@ -224,7 +224,7 @@ class RandomFlip3D(RandomFlip):
@TRANSFORMS.register_module()
class RandomJitterPoints(object):
class RandomJitterPoints(BaseTransform):
"""Randomly jitter point coordinates.
Different from the global translation in ``GlobalRotScaleTrans``, here we
......@@ -246,8 +246,8 @@ class RandomJitterPoints(object):
"""
def __init__(self,
jitter_std=[0.01, 0.01, 0.01],
clip_range=[-0.05, 0.05]):
jitter_std: List[float] = [0.01, 0.01, 0.01],
clip_range: List[float] = [-0.05, 0.05]) -> None:
seq_types = (list, tuple, np.ndarray)
if not isinstance(jitter_std, seq_types):
assert isinstance(jitter_std, (int, float)), \
......@@ -262,7 +262,7 @@ class RandomJitterPoints(object):
clip_range = [-clip_range, clip_range]
self.clip_range = clip_range
def __call__(self, input_dict):
def transform(self, input_dict: dict) -> dict:
"""Call function to jitter all the points in the scene.
Args:
......@@ -780,10 +780,10 @@ class GlobalRotScaleTrans(BaseTransform):
@TRANSFORMS.register_module()
class PointShuffle(object):
class PointShuffle(BaseTransform):
"""Shuffle input points."""
def __call__(self, input_dict):
def transform(self, input_dict: dict) -> dict:
"""Call function to shuffle points.
Args:
......@@ -1113,7 +1113,7 @@ class IndoorPointSample(PointSample):
@TRANSFORMS.register_module()
class IndoorPatchPointSample(object):
class IndoorPatchPointSample(BaseTransform):
r"""Indoor point sample within a patch. Modified from `PointNet++ <https://
github.com/charlesq34/pointnet2/blob/master/scannet/scannet_dataset.py>`_.
......@@ -1152,15 +1152,15 @@ class IndoorPatchPointSample(object):
"""
def __init__(self,
num_points,
block_size=1.5,
sample_rate=None,
ignore_index=None,
use_normalized_coord=False,
num_try=10,
enlarge_size=0.2,
min_unique_num=None,
eps=1e-2):
num_points: int,
block_size: float = 1.5,
sample_rate: Optional[float] = None,
ignore_index: Optional[int] = None,
use_normalized_coord: bool = False,
num_try: int = 10,
enlarge_size: float = 0.2,
min_unique_num: Optional[int] = None,
eps: float = 1e-2) -> None:
self.num_points = num_points
self.block_size = block_size
self.ignore_index = ignore_index
......@@ -1175,8 +1175,10 @@ class IndoorPatchPointSample(object):
"'sample_rate' has been deprecated and will be removed in "
'the future. Please remove them from your code.')
def _input_generation(self, coords, patch_center, coord_max, attributes,
attribute_dims, point_type):
def _input_generation(self, coords: np.ndarray, patch_center: np.ndarray,
coord_max: np.ndarray, attributes: np.ndarray,
attribute_dims: dict,
point_type: type) -> BasePoints:
"""Generating model input.
Generate input by subtracting patch center and adding additional
......@@ -1216,7 +1218,8 @@ class IndoorPatchPointSample(object):
return points
def _patch_points_sampling(self, points, sem_mask):
def _patch_points_sampling(self, points: BasePoints,
sem_mask: np.ndarray) -> BasePoints:
"""Patch points sampling.
First sample a valid patch.
......@@ -1316,7 +1319,7 @@ class IndoorPatchPointSample(object):
return points, choices
def __call__(self, results):
def transform(self, input_dict: dict) -> dict:
"""Call function to sample points to in indoor scenes.
Args:
......@@ -1326,22 +1329,33 @@ class IndoorPatchPointSample(object):
dict: Results after sampling, 'points', 'pts_instance_mask'
and 'pts_semantic_mask' keys are updated in the result dict.
"""
points = results['points']
points = input_dict['points']
assert 'pts_semantic_mask' in results.keys(), \
assert 'pts_semantic_mask' in input_dict.keys(), \
'semantic mask should be provided in training and evaluation'
pts_semantic_mask = results['pts_semantic_mask']
pts_semantic_mask = input_dict['pts_semantic_mask']
points, choices = self._patch_points_sampling(points,
pts_semantic_mask)
results['points'] = points
results['pts_semantic_mask'] = pts_semantic_mask[choices]
pts_instance_mask = results.get('pts_instance_mask', None)
input_dict['points'] = points
input_dict['pts_semantic_mask'] = pts_semantic_mask[choices]
# 'eval_ann_info' will be passed to evaluator
if 'eval_ann_info' in input_dict:
input_dict['eval_ann_info']['pts_semantic_mask'] = \
pts_semantic_mask[choices]
pts_instance_mask = input_dict.get('pts_instance_mask', None)
if pts_instance_mask is not None:
results['pts_instance_mask'] = pts_instance_mask[choices]
input_dict['pts_instance_mask'] = pts_instance_mask[choices]
# 'eval_ann_info' will be passed to evaluator
if 'eval_ann_info' in input_dict:
input_dict['eval_ann_info']['pts_instance_mask'] = \
pts_instance_mask[choices]
return results
return input_dict
def __repr__(self):
"""str: Return a string that describes the module."""
......@@ -1358,14 +1372,14 @@ class IndoorPatchPointSample(object):
@TRANSFORMS.register_module()
class BackgroundPointsFilter(object):
class BackgroundPointsFilter(BaseTransform):
"""Filter background points near the bounding box.
Args:
bbox_enlarge_range (tuple[float], float): Bbox enlarge range.
"""
def __init__(self, bbox_enlarge_range):
def __init__(self, bbox_enlarge_range: Union[Tuple[float], float]) -> None:
assert (is_tuple_of(bbox_enlarge_range, float)
and len(bbox_enlarge_range) == 3) \
or isinstance(bbox_enlarge_range, float), \
......@@ -1376,7 +1390,7 @@ class BackgroundPointsFilter(object):
self.bbox_enlarge_range = np.array(
bbox_enlarge_range, dtype=np.float32)[np.newaxis, :]
def __call__(self, input_dict):
def transform(self, input_dict: dict) -> dict:
"""Call function to filter points by the range.
Args:
......
mmdet3d/datasets/s3dis_dataset.py
View file @ 360c27f9
# Copyright (c) OpenMMLab. All rights reserved.
from os import path as osp
from typing import Callable, List, Optional, Union
import numpy as np
from mmdet3d.core import show_seg_result
from mmdet3d.core.bbox import DepthInstance3DBoxes
from mmdet3d.registry import DATASETS
from .custom_3d_seg import Custom3DSegDataset
from .det3d_dataset import Det3DDataset
from .pipelines import Compose
from .seg3d_dataset import Seg3DDataset
@DATASETS.register_module()
......@@ -153,7 +153,7 @@ class S3DISDataset(Det3DDataset):
return Compose(pipeline)
class _S3DISSegDataset(Custom3DSegDataset):
class _S3DISSegDataset(Seg3DDataset):
r"""S3DIS Dataset for Semantic Segmentation Task.
This class is the inner dataset for S3DIS. Since S3DIS has 6 areas, we
......@@ -185,114 +185,44 @@ class _S3DISSegDataset(Custom3DSegDataset):
data. For scenes with many points, we may sample it several times.
Defaults to None.
"""
CLASSES = ('ceiling', 'floor', 'wall', 'beam', 'column', 'window', 'door',
'table', 'chair', 'sofa', 'bookcase', 'board', 'clutter')
VALID_CLASS_IDS = tuple(range(13))
ALL_CLASS_IDS = tuple(range(14)) # possibly with 'stair' class
PALETTE = [[0, 255, 0], [0, 0, 255], [0, 255, 255], [255, 255, 0],
[255, 0, 255], [100, 100, 255], [200, 200, 100],
[170, 120, 200], [255, 0, 0], [200, 100, 100], [10, 200, 100],
[200, 200, 200], [50, 50, 50]]
METAINFO = {
'CLASSES':
('ceiling', 'floor', 'wall', 'beam', 'column', 'window', 'door',
'table', 'chair', 'sofa', 'bookcase', 'board', 'clutter'),
'PALETTE': [[0, 255, 0], [0, 0, 255], [0, 255, 255], [255, 255, 0],
[255, 0, 255], [100, 100, 255], [200, 200, 100],
[170, 120, 200], [255, 0, 0], [200, 100, 100],
[10, 200, 100], [200, 200, 200], [50, 50, 50]],
'valid_class_ids':
tuple(range(13)),
'all_class_ids':
tuple(range(14)) # possibly with 'stair' class
}
def __init__(self,
data_root,
ann_file,
pipeline=None,
classes=None,
palette=None,
modality=None,
test_mode=False,
data_root: Optional[str] = None,
ann_file: str = '',
metainfo: Optional[dict] = None,
data_prefix: dict = dict(
pts='points', img='', instance_mask='', semantic_mask=''),
pipeline: List[Union[dict, Callable]] = [],
modality: dict = dict(use_lidar=True, use_camera=False),
ignore_index=None,
scene_idxs=None,
**kwargs):
test_mode=False,
**kwargs) -> None:
super().__init__(
data_root=data_root,
ann_file=ann_file,
metainfo=metainfo,
data_prefix=data_prefix,
pipeline=pipeline,
classes=classes,
palette=palette,
modality=modality,
test_mode=test_mode,
ignore_index=ignore_index,
scene_idxs=scene_idxs,
test_mode=test_mode,
**kwargs)
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:
- 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]
pts_semantic_mask_path = osp.join(self.data_root,
info['pts_semantic_mask_path'])
anns_results = dict(pts_semantic_mask_path=pts_semantic_mask_path)
return anns_results
def _build_default_pipeline(self):
"""Build the default pipeline for this dataset."""
pipeline = [
dict(
type='LoadPointsFromFile',
coord_type='DEPTH',
shift_height=False,
use_color=True,
load_dim=6,
use_dim=[0, 1, 2, 3, 4, 5]),
dict(
type='LoadAnnotations3D',
with_bbox_3d=False,
with_label_3d=False,
with_mask_3d=False,
with_seg_3d=True),
dict(
type='PointSegClassMapping',
valid_cat_ids=self.VALID_CLASS_IDS,
max_cat_id=np.max(self.ALL_CLASS_IDS)),
dict(
type='DefaultFormatBundle3D',
with_label=False,
class_names=self.CLASSES),
dict(type='Collect3D', keys=['points', 'pts_semantic_mask'])
]
return Compose(pipeline)
def show(self, results, out_dir, show=True, pipeline=None):
"""Results visualization.
Args:
results (list[dict]): List of bounding boxes results.
out_dir (str): Output directory of visualization result.
show (bool): Visualize the results online.
pipeline (list[dict], optional): raw data loading for showing.
Default: None.
"""
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']
file_name = osp.split(pts_path)[-1].split('.')[0]
points, gt_sem_mask = self._extract_data(
i, pipeline, ['points', 'pts_semantic_mask'], load_annos=True)
points = points.numpy()
pred_sem_mask = result['semantic_mask'].numpy()
show_seg_result(points, gt_sem_mask,
pred_sem_mask, out_dir, file_name,
np.array(self.PALETTE), self.ignore_index, show)
def get_scene_idxs(self, scene_idxs):
"""Compute scene_idxs for data sampling.
......@@ -341,16 +271,17 @@ class S3DISSegDataset(_S3DISSegDataset):
"""
def __init__(self,
data_root,
ann_files,
pipeline=None,
classes=None,
palette=None,
modality=None,
test_mode=False,
data_root: Optional[str] = None,
ann_files: str = '',
metainfo: Optional[dict] = None,
data_prefix: dict = dict(
pts='points', img='', instance_mask='', semantic_mask=''),
pipeline: List[Union[dict, Callable]] = [],
modality: dict = dict(use_lidar=True, use_camera=False),
ignore_index=None,
scene_idxs=None,
**kwargs):
test_mode=False,
**kwargs) -> None:
# make sure that ann_files and scene_idxs have same length
ann_files = self._check_ann_files(ann_files)
......@@ -360,45 +291,45 @@ class S3DISSegDataset(_S3DISSegDataset):
super().__init__(
data_root=data_root,
ann_file=ann_files[0],
metainfo=metainfo,
data_prefix=data_prefix,
pipeline=pipeline,
classes=classes,
palette=palette,
modality=modality,
test_mode=test_mode,
ignore_index=ignore_index,
scene_idxs=scene_idxs[0],
test_mode=test_mode,
**kwargs)
datasets = [
_S3DISSegDataset(
data_root=data_root,
ann_file=ann_files[i],
metainfo=metainfo,
data_prefix=data_prefix,
pipeline=pipeline,
classes=classes,
palette=palette,
modality=modality,
test_mode=test_mode,
ignore_index=ignore_index,
scene_idxs=scene_idxs[i],
test_mode=test_mode,
**kwargs) for i in range(len(ann_files))
]
# data_infos and scene_idxs need to be concat
self.concat_data_infos([dst.data_infos for dst in datasets])
# data_list and scene_idxs need to be concat
self.concat_data_list([dst.data_list for dst in datasets])
self.concat_scene_idxs([dst.scene_idxs for dst in datasets])
# set group flag for the sampler
if not self.test_mode:
self._set_group_flag()
def concat_data_infos(self, data_infos):
"""Concat data_infos from several datasets to form self.data_infos.
def concat_data_list(self, data_lists):
"""Concat data_list from several datasets to form self.data_list.
Args:
data_infos (list[list[dict]])
data_lists (list[list[dict]])
"""
self.data_infos = [
info for one_data_infos in data_infos for info in one_data_infos
self.data_list = [
data for data_list in data_lists for data in data_list
]
def concat_scene_idxs(self, scene_idxs):
......
mmdet3d/datasets/scannet_dataset.py
View file @ 360c27f9
# Copyright (c) OpenMMLab. All rights reserved.
import tempfile
import warnings
from os import path as osp
from typing import Callable, List, Union
from typing import Callable, List, Optional, Union
import numpy as np
from mmdet3d.core import instance_seg_eval, show_result, show_seg_result
from mmdet3d.core import show_result
from mmdet3d.core.bbox import DepthInstance3DBoxes
from mmdet3d.registry import DATASETS
from .custom_3d_seg import Custom3DSegDataset
from .det3d_dataset import Det3DDataset
from .pipelines import Compose
from .seg3d_dataset import Seg3DDataset
@DATASETS.register_module()
......@@ -193,7 +192,7 @@ class ScanNetDataset(Det3DDataset):
@DATASETS.register_module()
class ScanNetSegDataset(Custom3DSegDataset):
class ScanNetSegDataset(Seg3DDataset):
r"""ScanNet Dataset for Semantic Segmentation Task.
This class serves as the API for experiments on the ScanNet Dataset.
......@@ -221,135 +220,64 @@ class ScanNetSegDataset(Custom3DSegDataset):
data. For scenes with many points, we may sample it several times.
Defaults to None.
"""
CLASSES = ('wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table',
'door', 'window', 'bookshelf', 'picture', 'counter', 'desk',
'curtain', 'refrigerator', 'showercurtrain', 'toilet', 'sink',
'bathtub', 'otherfurniture')
VALID_CLASS_IDS = (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28,
33, 34, 36, 39)
ALL_CLASS_IDS = tuple(range(41))
PALETTE = [
[174, 199, 232],
[152, 223, 138],
[31, 119, 180],
[255, 187, 120],
[188, 189, 34],
[140, 86, 75],
[255, 152, 150],
[214, 39, 40],
[197, 176, 213],
[148, 103, 189],
[196, 156, 148],
[23, 190, 207],
[247, 182, 210],
[219, 219, 141],
[255, 127, 14],
[158, 218, 229],
[44, 160, 44],
[112, 128, 144],
[227, 119, 194],
[82, 84, 163],
]
METAINFO = {
'CLASSES':
('wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table', 'door',
'window', 'bookshelf', 'picture', 'counter', 'desk', 'curtain',
'refrigerator', 'showercurtrain', 'toilet', 'sink', 'bathtub',
'otherfurniture'),
'PALETTE': [
[174, 199, 232],
[152, 223, 138],
[31, 119, 180],
[255, 187, 120],
[188, 189, 34],
[140, 86, 75],
[255, 152, 150],
[214, 39, 40],
[197, 176, 213],
[148, 103, 189],
[196, 156, 148],
[23, 190, 207],
[247, 182, 210],
[219, 219, 141],
[255, 127, 14],
[158, 218, 229],
[44, 160, 44],
[112, 128, 144],
[227, 119, 194],
[82, 84, 163],
],
'valid_class_ids': (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24,
28, 33, 34, 36, 39),
'all_class_ids':
tuple(range(41)),
}
def __init__(self,
data_root,
ann_file,
pipeline=None,
classes=None,
palette=None,
modality=None,
test_mode=False,
data_root: Optional[str] = None,
ann_file: str = '',
metainfo: Optional[dict] = None,
data_prefix: dict = dict(
pts='points', img='', instance_mask='', semantic_mask=''),
pipeline: List[Union[dict, Callable]] = [],
modality: dict = dict(use_lidar=True, use_camera=False),
ignore_index=None,
scene_idxs=None,
**kwargs):
test_mode=False,
**kwargs) -> None:
super().__init__(
data_root=data_root,
ann_file=ann_file,
metainfo=metainfo,
data_prefix=data_prefix,
pipeline=pipeline,
classes=classes,
palette=palette,
modality=modality,
test_mode=test_mode,
ignore_index=ignore_index,
scene_idxs=scene_idxs,
test_mode=test_mode,
**kwargs)
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:
- 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]
pts_semantic_mask_path = osp.join(self.data_root,
info['pts_semantic_mask_path'])
anns_results = dict(pts_semantic_mask_path=pts_semantic_mask_path)
return anns_results
def _build_default_pipeline(self):
"""Build the default pipeline for this dataset."""
pipeline = [
dict(
type='LoadPointsFromFile',
coord_type='DEPTH',
shift_height=False,
use_color=True,
load_dim=6,
use_dim=[0, 1, 2, 3, 4, 5]),
dict(
type='LoadAnnotations3D',
with_bbox_3d=False,
with_label_3d=False,
with_mask_3d=False,
with_seg_3d=True),
dict(
type='PointSegClassMapping',
valid_cat_ids=self.VALID_CLASS_IDS,
max_cat_id=np.max(self.ALL_CLASS_IDS)),
dict(
type='DefaultFormatBundle3D',
with_label=False,
class_names=self.CLASSES),
dict(type='Collect3D', keys=['points', 'pts_semantic_mask'])
]
return Compose(pipeline)
def show(self, results, out_dir, show=True, pipeline=None):
"""Results visualization.
Args:
results (list[dict]): List of bounding boxes results.
out_dir (str): Output directory of visualization result.
show (bool): Visualize the results online.
pipeline (list[dict], optional): raw data loading for showing.
Default: None.
"""
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']
file_name = osp.split(pts_path)[-1].split('.')[0]
points, gt_sem_mask = self._extract_data(
i, pipeline, ['points', 'pts_semantic_mask'], load_annos=True)
points = points.numpy()
pred_sem_mask = result['semantic_mask'].numpy()
show_seg_result(points, gt_sem_mask,
pred_sem_mask, out_dir, file_name,
np.array(self.PALETTE), self.ignore_index, show)
def get_scene_idxs(self, scene_idxs):
"""Compute scene_idxs for data sampling.
......@@ -362,191 +290,65 @@ class ScanNetSegDataset(Custom3DSegDataset):
return super().get_scene_idxs(scene_idxs)
def format_results(self, results, txtfile_prefix=None):
r"""Format the results to txt file. Refer to `ScanNet documentation
<http://kaldir.vc.in.tum.de/scannet_benchmark/documentation>`_.
Args:
outputs (list[dict]): Testing results of the dataset.
txtfile_prefix (str): The prefix of saved 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: (outputs, tmp_dir), outputs is the detection results,
tmp_dir is the temporal directory created for saving submission
files when ``submission_prefix`` is not specified.
"""
import mmcv
if txtfile_prefix is None:
tmp_dir = tempfile.TemporaryDirectory()
txtfile_prefix = osp.join(tmp_dir.name, 'results')
else:
tmp_dir = None
mmcv.mkdir_or_exist(txtfile_prefix)
# need to map network output to original label idx
pred2label = np.zeros(len(self.VALID_CLASS_IDS)).astype(np.int)
for original_label, output_idx in self.label_map.items():
if output_idx != self.ignore_index:
pred2label[output_idx] = original_label
outputs = []
for i, result in enumerate(results):
info = self.data_infos[i]
sample_idx = info['point_cloud']['lidar_idx']
pred_sem_mask = result['semantic_mask'].numpy().astype(np.int)
pred_label = pred2label[pred_sem_mask]
curr_file = f'{txtfile_prefix}/{sample_idx}.txt'
np.savetxt(curr_file, pred_label, fmt='%d')
outputs.append(dict(seg_mask=pred_label))
return outputs, tmp_dir
@DATASETS.register_module()
class ScanNetInstanceSegDataset(Custom3DSegDataset):
CLASSES = ('cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window',
'bookshelf', 'picture', 'counter', 'desk', 'curtain',
'refrigerator', 'showercurtrain', 'toilet', 'sink', 'bathtub',
'garbagebin')
VALID_CLASS_IDS = (3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34,
36, 39)
ALL_CLASS_IDS = tuple(range(41))
def get_ann_info(self, index):
"""Get annotation info according to the given index.
Args:
index (int): Index of the annotation data to get.
class ScanNetInstanceSegDataset(Seg3DDataset):
Returns:
dict: annotation information consists of the following keys:
- pts_semantic_mask_path (str): Path of semantic masks.
- pts_instance_mask_path (str): Path of instance masks.
"""
# Use index to get the annos, thus the evalhook could also use this api
info = self.data_infos[index]
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'])
anns_results = dict(
pts_instance_mask_path=pts_instance_mask_path,
pts_semantic_mask_path=pts_semantic_mask_path)
return anns_results
def get_classes_and_palette(self, classes=None, palette=None):
"""Get class names of current dataset. Palette is simply ignored for
instance segmentation.
Args:
classes (Sequence[str] | str | None): If classes is None, use
default CLASSES defined by builtin dataset. If classes is a
string, take it as a file name. The file contains the name of
classes where each line contains one class name. If classes is
a tuple or list, override the CLASSES defined by the dataset.
Defaults to None.
palette (Sequence[Sequence[int]]] | np.ndarray | None):
The palette of segmentation map. If None is given, random
palette will be generated. Defaults to None.
"""
if classes is not None:
return classes, None
return self.CLASSES, None
def _build_default_pipeline(self):
"""Build the default pipeline for this dataset."""
pipeline = [
dict(
type='LoadPointsFromFile',
coord_type='DEPTH',
shift_height=False,
use_color=True,
load_dim=6,
use_dim=[0, 1, 2, 3, 4, 5]),
dict(
type='LoadAnnotations3D',
with_bbox_3d=False,
with_label_3d=False,
with_mask_3d=True,
with_seg_3d=True),
dict(
type='PointSegClassMapping',
valid_cat_ids=self.VALID_CLASS_IDS,
max_cat_id=40),
dict(
type='DefaultFormatBundle3D',
with_label=False,
class_names=self.CLASSES),
dict(
type='Collect3D',
keys=['points', 'pts_semantic_mask', 'pts_instance_mask'])
]
return Compose(pipeline)
def evaluate(self,
results,
metric=None,
options=None,
logger=None,
show=False,
out_dir=None,
pipeline=None):
"""Evaluation in instance segmentation protocol.
Args:
results (list[dict]): List of results.
metric (str | list[str]): Metrics to be evaluated.
options (dict, optional): options for instance_seg_eval.
logger (logging.Logger | None | str): Logger used for printing
related information during evaluation. Defaults to None.
show (bool, optional): Whether to visualize.
Defaults to False.
out_dir (str, optional): Path to save the visualization results.
Defaults to None.
pipeline (list[dict], optional): raw data loading for showing.
Default: None.
METAINFO = {
'CLASSES':
('cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window',
'bookshelf', 'picture', 'counter', 'desk', 'curtain', 'refrigerator',
'showercurtrain', 'toilet', 'sink', 'bathtub', 'garbagebin'),
'PLATTE': [
[174, 199, 232],
[152, 223, 138],
[31, 119, 180],
[255, 187, 120],
[188, 189, 34],
[140, 86, 75],
[255, 152, 150],
[214, 39, 40],
[197, 176, 213],
[148, 103, 189],
[196, 156, 148],
[23, 190, 207],
[247, 182, 210],
[219, 219, 141],
[255, 127, 14],
[158, 218, 229],
[44, 160, 44],
[112, 128, 144],
[227, 119, 194],
[82, 84, 163],
],
'valid_class_ids':
(3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39),
'all_class_ids':
tuple(range(41))
}
Returns:
dict: Evaluation results.
"""
assert isinstance(
results, list), f'Expect results to be list, got {type(results)}.'
assert len(results) > 0, 'Expect length of results > 0.'
assert len(results) == len(self.data_infos)
assert isinstance(
results[0], dict
), f'Expect elements in results to be dict, got {type(results[0])}.'
load_pipeline = self._get_pipeline(pipeline)
pred_instance_masks = [result['instance_mask'] for result in results]
pred_instance_labels = [result['instance_label'] for result in results]
pred_instance_scores = [result['instance_score'] for result in results]
gt_semantic_masks, gt_instance_masks = zip(*[
self._extract_data(
index=i,
pipeline=load_pipeline,
key=['pts_semantic_mask', 'pts_instance_mask'],
load_annos=True) for i in range(len(self.data_infos))
])
ret_dict = instance_seg_eval(
gt_semantic_masks,
gt_instance_masks,
pred_instance_masks,
pred_instance_labels,
pred_instance_scores,
valid_class_ids=self.VALID_CLASS_IDS,
class_labels=self.CLASSES,
options=options,
logger=logger)
if show:
raise NotImplementedError('show is not implemented for now')
return ret_dict
def __init__(self,
data_root: Optional[str] = None,
ann_file: str = '',
metainfo: Optional[dict] = None,
data_prefix: dict = dict(
pts='points', img='', instance_mask='', semantic_mask=''),
pipeline: List[Union[dict, Callable]] = [],
modality: dict = dict(use_lidar=True, use_camera=False),
test_mode=False,
ignore_index=None,
scene_idxs=None,
file_client_args=dict(backend='disk'),
**kwargs) -> None:
super().__init__(
data_root=data_root,
ann_file=ann_file,
metainfo=metainfo,
pipeline=pipeline,
data_prefix=data_prefix,
modality=modality,
test_mode=test_mode,
ignore_index=ignore_index,
scene_idxs=scene_idxs,
file_client_args=file_client_args,
**kwargs)
mmdet3d/datasets/seg3d_dataset.py 0 → 100644
View file @ 360c27f9
# Copyright (c) OpenMMLab. All rights reserved.
from os import path as osp
from typing import Callable, Dict, List, Optional, Sequence, Union
import mmcv
import numpy as np
from mmengine.dataset import BaseDataset
from mmdet3d.registry import DATASETS
@DATASETS.register_module()
class Seg3DDataset(BaseDataset):
"""Base Class for 3D semantic segmentation dataset.
This is the base dataset of ScanNet, S3DIS and SemanticKITTI dataset.
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.
metainfo (dict, optional): Meta information for dataset, such as class
information. Defaults to None.
data_prefix (dict, optional): Prefix for training data. Defaults to
dict(pts='velodyne', img='', instance_mask='', semantic_mask='').
pipeline (list[dict], optional): Pipeline used for data processing.
Defaults to None.
modality (dict, optional): Modality to specify the sensor data used
as input, it usually has following keys.
- use_camera: bool
- use_lidar: bool
Defaults to `dict(use_lidar=True, use_camera=False)`
test_mode (bool, optional): Whether the dataset is in test mode.
Defaults to False.
ignore_index (int, optional): The label index to be ignored, e.g.
unannotated points. If None is given, set to len(self.CLASSES) to
be consistent with PointSegClassMapping function in pipeline.
Defaults to None.
scene_idxs (np.ndarray | str, optional): Precomputed index to load
data. For scenes with many points, we may sample it several times.
Defaults to None.
load_eval_anns (bool): Whether to load annotations
in test_mode, the annotation will be save in
`eval_ann_infos`, which can be use in Evaluator.
file_client_args (dict): Configuration of file client.
Defaults to `dict(backend='disk')`.
"""
METAINFO = {
'CLASSES': None, # names of all classes data used for the task
'PALETTE': None, # official color for visualization
'valid_class_ids': None, # class_ids used for training
'all_class_ids': None, # all possible class_ids in loaded seg mask
}
def __init__(self,
data_root: Optional[str] = None,
ann_file: str = '',
metainfo: Optional[dict] = None,
data_prefix: dict = dict(
pts='points',
img='',
pts_instance_mask='',
pts_emantic_mask=''),
pipeline: List[Union[dict, Callable]] = [],
modality: dict = dict(use_lidar=True, use_camera=False),
ignore_index: Optional[int] = None,
scene_idxs: Optional[str] = None,
test_mode: bool = False,
load_eval_anns: bool = True,
file_client_args: dict = dict(backend='disk'),
**kwargs) -> None:
# init file client
self.file_client = mmcv.FileClient(**file_client_args)
self.modality = modality
self.load_eval_anns = load_eval_anns
# TODO: We maintain the ignore_index attributes,
# but we may consider to remove it in the future.
self.ignore_index = len(self.METAINFO['CLASSES']) if \
ignore_index is None else ignore_index
# Get label mapping for custom classes
new_classes = metainfo.get('CLASSES', None)
self.label_mapping, self.label2cat, valid_class_ids = \
self.get_label_mapping(new_classes)
metainfo['label_mapping'] = self.label_mapping
metainfo['label2cat'] = self.label2cat
metainfo['valid_class_ids'] = valid_class_ids
# generate palette if it is not defined based on
# label mapping, otherwise directly use palette
# defined in dataset config.
palette = metainfo.get('PALETTE', None)
updated_palette = self._update_palette(new_classes, palette)
metainfo['PALETTE'] = updated_palette
super().__init__(
ann_file=ann_file,
metainfo=metainfo,
data_root=data_root,
data_prefix=data_prefix,
pipeline=pipeline,
test_mode=test_mode,
**kwargs)
self.scene_idxs = self.get_scene_idxs(scene_idxs)
# set group flag for the sampler
if not self.test_mode:
self._set_group_flag()
def get_label_mapping(self,
new_classes: Optional[Sequence] = None
) -> Union[Dict, None]:
"""Get label mapping.
The ``label_mapping`` is a dictionary, its keys are the old label ids
and its values are the new label ids, and is used for changing pixel
labels in load_annotations. If and only if old classes in cls.METAINFO
is not equal to new classes in self._metainfo and nether of them is not
None, `label_mapping` is not None.
Args:
new_classes (list, tuple, optional): The new classes name from
metainfo. Default to None.
Returns:
tuple: The mapping from old classes in cls.METAINFO to
new classes in metainfo
"""
old_classes = self.METAINFO.get('CLASSSES', None)
if (new_classes is not None and old_classes is not None
and list(new_classes) != list(old_classes)):
label_mapping = {}
if not set(new_classes).issubset(old_classes):
raise ValueError(
f'new classes {new_classes} is not a '
f'subset of CLASSES {old_classes} in METAINFO.')
# obtain true id from valid_class_ids
valid_class_ids = [
self.METAINFO['valid_class_ids'][old_classes.index(cls_name)]
for cls_name in new_classes
]
label_mapping = {
cls_id: self.ignore_index
for cls_id in self.METAINFO['all_class_ids']
}
label_mapping.update(
{cls_id: i
for i, cls_id in enumerate(valid_class_ids)})
label2cat = {i: cat_name for i, cat_name in enumerate(new_classes)}
else:
label_mapping = {
cls_id: self.ignore_index
for cls_id in self.METAINFO['all_class_ids']
}
label_mapping.update({
cls_id: i
for i, cls_id in enumerate(self.METAINFO['valid_class_ids'])
})
# map label to category name
label2cat = {
i: cat_name
for i, cat_name in enumerate(self.METAINFO['CLASSES'])
}
valid_class_ids = self.METAINFO['valid_class_ids']
return label_mapping, label2cat, valid_class_ids
def _update_palette(self, new_classes, palette) -> list:
"""Update palette according to metainfo.
If length of palette is equal to classes, just return the palette.
If palette is not defined, it will randomly generate a palette.
If classes is updated by customer, it will return the subset of
palette.
Returns:
Sequence: Palette for current dataset.
"""
if palette is None:
# If palette is not defined, it generate a palette according
# to the original PALETTE and classes.
old_classes = self.METAINFO.get('CLASSSES', None)
palette = [
self.METAINFO['PALETTE'][old_classes.index(cls_name)]
for cls_name in new_classes
]
return palette
# palette does match classes
if len(palette) == len(new_classes):
return palette
else:
raise ValueError('Once PLATTE in set in metainfo, it should'
'match CLASSES in metainfo')
def parse_data_info(self, info: dict) -> dict:
"""Process the raw data info.
Convert all relative path of needed modality data file to
the absolute path. And process
the `instances` field to `ann_info` in training stage.
Args:
info (dict): Raw info dict.
Returns:
dict: Has `ann_info` in training stage. And
all path has been converted to absolute path.
"""
if self.modality['use_lidar']:
info['lidar_points']['lidar_path'] = \
osp.join(
self.data_prefix.get('pts', ''),
info['lidar_points']['lidar_path'])
if self.modality['use_camera']:
for cam_id, img_info in info['images'].items():
if 'img_path' in img_info:
img_info['img_path'] = osp.join(
self.data_prefix.get('img', ''), img_info['img_path'])
if 'pts_instance_mask_path' in info:
info['pts_instance_mask_path'] = \
osp.join(self.data_prefix.get('pts_instance_mask', ''),
info['pts_instance_mask_path'])
if 'pts_semantic_mask_path' in info:
info['pts_semantic_mask_path'] = \
osp.join(self.data_prefix.get('pts_semantic_mask', ''),
info['pts_semantic_mask_path'])
# Add label_mapping to input dict for directly
# use it in PointSegClassMapping pipeline
info['label_mapping'] = self.label_mapping
# 'eval_ann_info' will be updated in loading pipelines
if self.test_mode and self.load_eval_anns:
info['eval_ann_info'] = dict()
return info
def get_scene_idxs(self, scene_idxs):
"""Compute scene_idxs for data sampling.
We sample more times for scenes with more points.
"""
if self.test_mode:
# when testing, we load one whole scene every time
return np.arange(len(self.data_list)).astype(np.int32)
# we may need to re-sample different scenes according to scene_idxs
# this is necessary for indoor scene segmentation such as ScanNet
if scene_idxs is None:
scene_idxs = np.arange(len(self.data_list))
if isinstance(scene_idxs, str):
with self.file_client.get_local_path(scene_idxs) as local_path:
scene_idxs = np.load(local_path)
else:
scene_idxs = np.array(scene_idxs)
return scene_idxs.astype(np.int32)
def _set_group_flag(self):
"""Set flag according to image aspect ratio.
Images with aspect ratio greater than 1 will be set as group 1,
otherwise group 0. In 3D datasets, they are all the same, thus are all
zeros.
"""
self.flag = np.zeros(len(self), dtype=np.uint8)
mmdet3d/datasets/semantickitti_dataset.py
View file @ 360c27f9
# Copyright (c) OpenMMLab. All rights reserved.
from os import path as osp
from typing import Callable, List, Optional, Union
from mmdet3d.registry import DATASETS
from .det3d_dataset import Det3DDataset
from .seg3d_dataset import Seg3DDataset
@DATASETS.register_module()
class SemanticKITTIDataset(Det3DDataset):
class SemanticKITTIDataset(Seg3DDataset):
r"""SemanticKITTI Dataset.
This class serves as the API for experiments on the SemanticKITTI Dataset
......@@ -36,75 +36,38 @@ class SemanticKITTIDataset(Det3DDataset):
test_mode (bool, optional): Whether the dataset is in test mode.
Defaults to False.
"""
CLASSES = ('unlabeled', 'car', 'bicycle', 'motorcycle', 'truck', 'bus',
'person', 'bicyclist', 'motorcyclist', 'road', 'parking',
'sidewalk', 'other-ground', 'building', 'fence', 'vegetation',
'trunck', 'terrian', 'pole', 'traffic-sign')
METAINFO = {
'CLASSES': ('unlabeled', 'car', 'bicycle', 'motorcycle', 'truck',
'bus', 'person', 'bicyclist', 'motorcyclist', 'road',
'parking', 'sidewalk', 'other-ground', 'building', 'fence',
'vegetation', 'trunck', 'terrian', 'pole', 'traffic-sign'),
'valid_class_ids':
tuple(range(20)),
'all_class_ids':
tuple(range(20))
}
def __init__(self,
data_root,
ann_file,
pipeline=None,
classes=None,
modality=None,
box_type_3d='Lidar',
filter_empty_gt=False,
test_mode=False):
data_root: Optional[str] = None,
ann_file: str = '',
metainfo: Optional[dict] = None,
data_prefix: dict = dict(
pts='points', img='', instance_mask='', semantic_mask=''),
pipeline: List[Union[dict, Callable]] = [],
modality: dict = dict(use_lidar=True, use_camera=False),
ignore_index=None,
scene_idxs=None,
test_mode=False,
**kwargs) -> None:
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)
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.
- file_name (str): Filename of point clouds.
- 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(
pts_filename=pts_filename,
sample_idx=sample_idx,
file_name=pts_filename)
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.
Args:
index (int): Index of the annotation data to get.
Returns:
dict: annotation information consists of the following keys:
- 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]
pts_semantic_mask_path = osp.join(self.data_root,
info['pts_semantic_mask_path'])
anns_results = dict(pts_semantic_mask_path=pts_semantic_mask_path)
return anns_results
ignore_index=ignore_index,
scene_idxs=scene_idxs,
test_mode=test_mode,
**kwargs)
tests/test_core/test_data_structure/test_det_data_sample.py
View file @ 360c27f9
# Copyright (c) OpenMMLab. All rights reserved.
from unittest import TestCase
import numpy as np
import pytest
import torch
from mmengine.data import InstanceData, PixelData
from mmengine.data import InstanceData
from mmdet3d.core.data_structures import Det3DDataSample
from mmdet3d.core.data_structures import Det3DDataSample, PointData
def _equal(a, b):
......@@ -86,47 +87,34 @@ class TestDet3DataSample(TestCase):
assert _equal(det3d_data_sample.img_pred_instances_3d.scores_3d,
img_pred_instances_3d_data['scores_3d'])
# test gt_panoptic_seg
gt_pts_panoptic_seg_data = dict(panoptic_seg=torch.rand(5, 4))
gt_pts_panoptic_seg = PixelData(**gt_pts_panoptic_seg_data)
det3d_data_sample.gt_pts_panoptic_seg = gt_pts_panoptic_seg
assert 'gt_pts_panoptic_seg' in det3d_data_sample
assert _equal(det3d_data_sample.gt_pts_panoptic_seg.panoptic_seg,
gt_pts_panoptic_seg_data['panoptic_seg'])
# test pred_panoptic_seg
pred_pts_panoptic_seg_data = dict(panoptic_seg=torch.rand(5, 4))
pred_pts_panoptic_seg = PixelData(**pred_pts_panoptic_seg_data)
det3d_data_sample.pred_pts_panoptic_seg = pred_pts_panoptic_seg
assert 'pred_pts_panoptic_seg' in det3d_data_sample
assert _equal(det3d_data_sample.pred_pts_panoptic_seg.panoptic_seg,
pred_pts_panoptic_seg_data['panoptic_seg'])
# test gt_sem_seg
gt_pts_sem_seg_data = dict(segm_seg=torch.rand(5, 4, 2))
gt_pts_sem_seg = PixelData(**gt_pts_sem_seg_data)
det3d_data_sample.gt_pts_sem_seg = gt_pts_sem_seg
assert 'gt_pts_sem_seg' in det3d_data_sample
assert _equal(det3d_data_sample.gt_pts_sem_seg.segm_seg,
gt_pts_sem_seg_data['segm_seg'])
# test pred_segm_seg
pred_pts_sem_seg_data = dict(segm_seg=torch.rand(5, 4, 2))
pred_pts_sem_seg = PixelData(**pred_pts_sem_seg_data)
det3d_data_sample.pred_pts_sem_seg = pred_pts_sem_seg
assert 'pred_pts_sem_seg' in det3d_data_sample
assert _equal(det3d_data_sample.pred_pts_sem_seg.segm_seg,
pred_pts_sem_seg_data['segm_seg'])
# test gt_seg
gt_pts_seg_data = dict(
pts_instance_mask=torch.rand(20), pts_semantic_mask=torch.rand(20))
gt_pts_seg = PointData(**gt_pts_seg_data)
det3d_data_sample.gt_pts_seg = gt_pts_seg
assert 'gt_pts_seg' in det3d_data_sample
assert _equal(det3d_data_sample.gt_pts_seg.pts_instance_mask,
gt_pts_seg_data['pts_instance_mask'])
assert _equal(det3d_data_sample.gt_pts_seg.pts_semantic_mask,
gt_pts_seg_data['pts_semantic_mask'])
# test pred_seg
pred_pts_seg_data = dict(
pts_instance_mask=torch.rand(20), pts_semantic_mask=torch.rand(20))
pred_pts_seg = PointData(**pred_pts_seg_data)
det3d_data_sample.pred_pts_seg = pred_pts_seg
assert 'pred_pts_seg' in det3d_data_sample
assert _equal(det3d_data_sample.pred_pts_seg.pts_instance_mask,
pred_pts_seg_data['pts_instance_mask'])
assert _equal(det3d_data_sample.pred_pts_seg.pts_semantic_mask,
pred_pts_seg_data['pts_semantic_mask'])
# test type error
with pytest.raises(AssertionError):
det3d_data_sample.pred_instances_3d = torch.rand(2, 4)
with pytest.raises(AssertionError):
det3d_data_sample.pred_pts_panoptic_seg = torch.rand(2, 4)
with pytest.raises(AssertionError):
det3d_data_sample.pred_pts_sem_seg = torch.rand(2, 4)
det3d_data_sample.pred_pts_seg = torch.rand(20)
def test_deleter(self):
tmp_instances_3d_data = dict(
......@@ -157,17 +145,10 @@ class TestDet3DataSample(TestCase):
del det3d_data_sample.img_pred_instances_3d
assert 'img_pred_instances_3d' not in det3d_data_sample
pred_pts_panoptic_seg_data = torch.rand(5, 4)
pred_pts_panoptic_seg_data = PixelData(data=pred_pts_panoptic_seg_data)
det3d_data_sample.pred_pts_panoptic_seg_data = \
pred_pts_panoptic_seg_data
assert 'pred_pts_panoptic_seg_data' in det3d_data_sample
del det3d_data_sample.pred_pts_panoptic_seg_data
assert 'pred_pts_panoptic_seg_data' not in det3d_data_sample
pred_pts_sem_seg_data = dict(segm_seg=torch.rand(5, 4, 2))
pred_pts_sem_seg = PixelData(**pred_pts_sem_seg_data)
det3d_data_sample.pred_pts_sem_seg = pred_pts_sem_seg
assert 'pred_pts_sem_seg' in det3d_data_sample
del det3d_data_sample.pred_pts_sem_seg
assert 'pred_pts_sem_seg' not in det3d_data_sample
pred_pts_seg_data = dict(
pts_instance_mask=torch.rand(20), pts_semantic_mask=torch.rand(20))
pred_pts_seg = PointData(**pred_pts_seg_data)
det3d_data_sample.pred_pts_seg = pred_pts_seg
assert 'pred_pts_seg' in det3d_data_sample
del det3d_data_sample.pred_pts_seg
assert 'pred_pts_seg' not in det3d_data_sample
tests/test_data/test_datasets/test_s3dis_dataset.py 0 → 100644
View file @ 360c27f9
# Copyright (c) OpenMMLab. All rights reserved.
import unittest
import numpy as np
import torch
from mmdet3d.datasets import S3DISSegDataset
from mmdet3d.utils import register_all_modules
def _generate_s3dis_seg_dataset_config():
data_root = './tests/data/s3dis/'
ann_file = 's3dis_infos.pkl'
classes = ('ceiling', 'floor', 'wall', 'beam', 'column', 'window', 'door',
'table', 'chair', 'sofa', 'bookcase', 'board', 'clutter')
palette = [[0, 255, 0], [0, 0, 255], [0, 255, 255], [255, 255, 0],
[255, 0, 255], [100, 100, 255], [200, 200, 100],
[170, 120, 200], [255, 0, 0], [200, 100, 100], [10, 200, 100],
[200, 200, 200], [50, 50, 50]]
scene_idxs = [0 for _ in range(20)]
modality = dict(use_lidar=True, use_camera=False)
pipeline = [
dict(
type='LoadPointsFromFile',
coord_type='DEPTH',
shift_height=False,
use_color=True,
load_dim=6,
use_dim=[0, 1, 2, 3, 4, 5]),
dict(
type='LoadAnnotations3D',
with_bbox_3d=False,
with_label_3d=False,
with_mask_3d=False,
with_seg_3d=True),
dict(type='PointSegClassMapping'),
dict(
type='IndoorPatchPointSample',
num_points=5,
block_size=1.0,
ignore_index=len(classes),
use_normalized_coord=True,
enlarge_size=0.2,
min_unique_num=None),
dict(type='NormalizePointsColor', color_mean=None),
dict(type='Pack3DDetInputs', keys=['points', 'pts_semantic_mask'])
]
data_prefix = dict(
pts='points',
pts_instance_mask='instance_mask',
pts_semantic_mask='semantic_mask')
return (data_root, ann_file, classes, palette, scene_idxs, data_prefix,
pipeline, modality)
class TestS3DISDataset(unittest.TestCase):
def test_s3dis_seg(self):
np.random.seed(0)
data_root, ann_file, classes, palette, scene_idxs, data_prefix, \
pipeline, modality, = _generate_s3dis_seg_dataset_config()
register_all_modules()
s3dis_seg_dataset = S3DISSegDataset(
data_root,
ann_file,
metainfo=dict(CLASSES=classes, PALETTE=palette),
data_prefix=data_prefix,
pipeline=pipeline,
modality=modality,
scene_idxs=scene_idxs)
input_dict = s3dis_seg_dataset.prepare_data(0)
points = input_dict['inputs']['points']
data_sample = input_dict['data_sample']
pts_semantic_mask = data_sample.gt_pts_seg.pts_semantic_mask
expected_points = torch.tensor([[
0.0000, 0.0000, 3.1720, 0.4706, 0.4431, 0.3725, 0.4624, 0.7502,
0.9543
],
[
0.2880, -0.5900, 0.0650, 0.3451,
0.3373, 0.3490, 0.5119, 0.5518,
0.0196
],
[
0.1570, 0.6000, 3.1700, 0.4941,
0.4667, 0.3569, 0.4893, 0.9519,
0.9537
],
[
-0.1320, 0.3950, 0.2720, 0.3216,
0.2863, 0.2275, 0.4397, 0.8830,
0.0818
],
[
-0.4860, -0.0640, 3.1710, 0.3843,
0.3725, 0.3059, 0.3789, 0.7286,
0.9540
]])
expected_pts_semantic_mask = np.array([0, 1, 0, 8, 0])
assert torch.allclose(points, expected_points, 1e-2)
self.assertTrue(
(pts_semantic_mask.numpy() == expected_pts_semantic_mask).all())
tests/test_data/test_datasets/test_scannet_dataset.py
View file @ 360c27f9
......@@ -6,7 +6,74 @@ import torch
from mmengine.testing import assert_allclose
from mmdet3d.core import DepthInstance3DBoxes
from mmdet3d.datasets import ScanNetDataset
from mmdet3d.datasets import ScanNetDataset, ScanNetSegDataset
from mmdet3d.utils import register_all_modules
def _generate_scannet_seg_dataset_config():
data_root = './tests/data/scannet/'
ann_file = 'scannet_infos.pkl'
classes = ('wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table',
'door', 'window', 'bookshelf', 'picture', 'counter', 'desk',
'curtain', 'refrigerator', 'showercurtrain', 'toilet', 'sink',
'bathtub', 'otherfurniture')
palette = [
[174, 199, 232],
[152, 223, 138],
[31, 119, 180],
[255, 187, 120],
[188, 189, 34],
[140, 86, 75],
[255, 152, 150],
[214, 39, 40],
[197, 176, 213],
[148, 103, 189],
[196, 156, 148],
[23, 190, 207],
[247, 182, 210],
[219, 219, 141],
[255, 127, 14],
[158, 218, 229],
[44, 160, 44],
[112, 128, 144],
[227, 119, 194],
[82, 84, 163],
]
scene_idxs = [0 for _ in range(20)]
modality = dict(use_lidar=True, use_camera=False)
pipeline = [
dict(
type='LoadPointsFromFile',
coord_type='DEPTH',
shift_height=False,
use_color=True,
load_dim=6,
use_dim=[0, 1, 2, 3, 4, 5]),
dict(
type='LoadAnnotations3D',
with_bbox_3d=False,
with_label_3d=False,
with_mask_3d=False,
with_seg_3d=True),
dict(type='PointSegClassMapping'),
dict(
type='IndoorPatchPointSample',
num_points=5,
block_size=1.5,
ignore_index=len(classes),
use_normalized_coord=True,
enlarge_size=0.2,
min_unique_num=None),
dict(type='NormalizePointsColor', color_mean=None),
dict(type='Pack3DDetInputs', keys=['points', 'pts_semantic_mask'])
]
data_prefix = dict(
pts='points',
pts_instance_mask='instance_mask',
pts_semantic_mask='semantic_mask')
return (data_root, ann_file, classes, palette, scene_idxs, data_prefix,
pipeline, modality)
def _generate_scannet_dataset_config():
......@@ -92,3 +159,54 @@ class TestScanNetDataset(unittest.TestCase):
# 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)
def test_scannet_seg(self):
np.random.seed(0)
data_root, ann_file, classes, palette, scene_idxs, data_prefix, \
pipeline, modality, = _generate_scannet_seg_dataset_config()
register_all_modules()
scannet_seg_dataset = ScanNetSegDataset(
data_root,
ann_file,
metainfo=dict(CLASSES=classes, PALETTE=palette),
data_prefix=data_prefix,
pipeline=pipeline,
modality=modality,
scene_idxs=scene_idxs)
input_dict = scannet_seg_dataset.prepare_data(0)
points = input_dict['inputs']['points']
data_sample = input_dict['data_sample']
pts_semantic_mask = data_sample.gt_pts_seg.pts_semantic_mask
expected_points = torch.tensor([[
0.0000, 0.0000, 1.2427, 0.6118, 0.5529, 0.4471, -0.6462, -1.0046,
0.4280
],
[
0.1553, -0.0074, 1.6077, 0.5882,
0.6157, 0.5569, -0.6001, -1.0068,
0.5537
],
[
0.1518, 0.6016, 0.6548, 0.1490,
0.1059, 0.0431, -0.6012, -0.8309,
0.2255
],
[
-0.7494, 0.1033, 0.6756, 0.5216,
0.4353, 0.3333, -0.8687, -0.9748,
0.2327
],
[
-0.6836, -0.0203, 0.5884, 0.5765,
0.5020, 0.4510, -0.8491, -1.0105,
0.2027
]])
expected_pts_semantic_mask = np.array([13, 13, 12, 2, 0])
assert torch.allclose(points, expected_points, 1e-2)
self.assertTrue(
(pts_semantic_mask.numpy() == expected_pts_semantic_mask).all())
tests/test_data/test_datasets/test_semantickitti_dataset.py 0 → 100644
View file @ 360c27f9
# Copyright (c) OpenMMLab. All rights reserved.
import unittest
import numpy as np
from mmdet3d.datasets import SemanticKITTIDataset
from mmdet3d.utils import register_all_modules
def _generate_semantickitti_dataset_config():
data_root = './tests/data/semantickitti/'
ann_file = 'semantickitti_infos.pkl'
classes = ('unlabeled', 'car', 'bicycle', 'motorcycle', 'truck', 'bus',
'person', 'bicyclist', 'motorcyclist', 'road', 'parking',
'sidewalk', 'other-ground', 'building', 'fence', 'vegetation',
'trunck', 'terrian', 'pole', 'traffic-sign')
palette = [
[174, 199, 232],
[152, 223, 138],
[31, 119, 180],
[255, 187, 120],
[188, 189, 34],
[140, 86, 75],
[255, 152, 150],
[214, 39, 40],
[197, 176, 213],
[148, 103, 189],
[196, 156, 148],
[23, 190, 207],
[247, 182, 210],
[219, 219, 141],
[255, 127, 14],
[158, 218, 229],
[44, 160, 44],
[112, 128, 144],
[227, 119, 194],
[82, 84, 163],
]
modality = dict(use_lidar=True, use_camera=False)
pipeline = [
dict(
type='LoadPointsFromFile',
coord_type='LIDAR',
shift_height=True,
load_dim=4,
use_dim=[0, 1, 2]),
dict(
type='LoadAnnotations3D',
with_bbox_3d=False,
with_label_3d=False,
with_mask_3d=False,
with_seg_3d=True,
seg_3d_dtype=np.int32),
dict(type='Pack3DDetInputs', keys=['points', 'pts_semantic_mask'])
]
data_prefix = dict(
pts='sequences/00/velodyne', pts_semantic_mask='sequences/00/labels')
return (data_root, ann_file, classes, palette, data_prefix, pipeline,
modality)
class TestSemanticKITTIDataset(unittest.TestCase):
def test_semantickitti(self):
np.random.seed(0)
data_root, ann_file, classes, palette, data_prefix, \
pipeline, modality, = _generate_semantickitti_dataset_config()
register_all_modules()
semantickitti_dataset = SemanticKITTIDataset(
data_root,
ann_file,
metainfo=dict(CLASSES=classes, PALETTE=palette),
data_prefix=data_prefix,
pipeline=pipeline,
modality=modality)
input_dict = semantickitti_dataset.prepare_data(0)
points = input_dict['inputs']['points']
data_sample = input_dict['data_sample']
pts_semantic_mask = data_sample.gt_pts_seg.pts_semantic_mask
self.assertEqual(points.shape[0], pts_semantic_mask.shape[0])
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