Skip to content

GitLab

Unverified Commit a0090aa1 authored by Ziyi Wu's avatar Ziyi Wu Committed by GitHub
Browse files

[Feature] Support S3DIS data pre-processing and dataset class (#433) 

* support S3DIS data download and pre-processing (to ScanNet format)

* add S3DIS data for unittest

* add S3DIS semseg dataset class and unittest

* add config file for S3DIS dataset

* add eval_pipeline to S3DIS dataset config file

* clean code for S3DIS pre-processing scripts

* reformat code

* fix small bugs

* resolve conflicts & modify show() to use pipeline

* fix small errors

* polish data pre-processing code

* add more comments about S3DIS dataset

* fix markdown lint error
parent 3c540f71
Hide whitespace changes
Inline Side-by-side
Showing with 1366 additions and 6 deletions
.gitignore
View file @ a0090aa1
...@@ -123,6 +123,7 @@ exps/ ...@@ -123,6 +123,7 @@ exps/
# demo # demo
*.jpg *.jpg
*.png *.png
/data/s3dis/Stanford3dDataset_v1.2_Aligned_Version/
/data/scannet/scans/ /data/scannet/scans/
/data/sunrgbd/OFFICIAL_SUNRGBD/ /data/sunrgbd/OFFICIAL_SUNRGBD/
*.obj *.obj
......
configs/_base_/datasets/s3dis_seg-3d-13class.py 0 → 100644
View file @ a0090aa1
# dataset settings
dataset_type = 'S3DISSegDataset'
data_root = './data/s3dis/'
class_names = ('ceiling', 'floor', 'wall', 'beam', 'column', 'window', 'door',
'table', 'chair', 'sofa', 'bookcase', 'board', 'clutter')
num_points = 4096
train_area = [1, 2, 3, 4, 6]
test_area = 5
train_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=tuple(range(len(class_names)))),
dict(
type='IndoorPatchPointSample',
num_points=num_points,
block_size=1.0,
sample_rate=1.0,
ignore_index=len(class_names),
use_normalized_coord=True),
dict(type='NormalizePointsColor', color_mean=None),
dict(type='DefaultFormatBundle3D', class_names=class_names),
dict(type='Collect3D', keys=['points', 'pts_semantic_mask'])
]
test_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='NormalizePointsColor', color_mean=None),
dict(type='DefaultFormatBundle3D', class_names=class_names),
dict(type='Collect3D', 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)
# we need to load gt seg_mask!
eval_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=tuple(range(len(class_names)))),
dict(
type='DefaultFormatBundle3D',
with_label=False,
class_names=class_names),
dict(type='Collect3D', keys=['points', 'pts_semantic_mask'])
]
data = dict(
samples_per_gpu=8,
workers_per_gpu=4,
# train on area 1, 2, 3, 4, 6
# test on area 5
train=dict(
type=dataset_type,
data_root=data_root,
ann_files=[
data_root + f's3dis_infos_Area_{i}.pkl' for i in train_area
],
pipeline=train_pipeline,
classes=class_names,
test_mode=False,
ignore_index=len(class_names),
scene_idxs=[
data_root + f'seg_info/Area_{i}_resampled_scene_idxs.npy'
for i in train_area
],
label_weight=[
data_root + f'seg_info/Area_{i}_label_weight.npy'
for i in train_area
]),
val=dict(
type=dataset_type,
data_root=data_root,
ann_file=data_root + f's3dis_infos_Area_{test_area}.pkl',
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 + f's3dis_infos_Area_{test_area}.pkl',
pipeline=test_pipeline,
classes=class_names,
test_mode=True,
ignore_index=len(class_names)))
evaluation = dict(pipeline=eval_pipeline)
data/s3dis/README.md 0 → 100644
View file @ a0090aa1
### Prepare S3DIS Data
We follow the procedure in [pointnet](https://github.com/charlesq34/pointnet).
1. Download S3DIS data by filling this [Google form](https://docs.google.com/forms/d/e/1FAIpQLScDimvNMCGhy_rmBA2gHfDu3naktRm6A8BPwAWWDv-Uhm6Shw/viewform?c=0&w=1). Download the ```Stanford3dDataset_v1.2_Aligned_Version.zip``` file and unzip it. Link or move the folder to this level of directory.
2. In this directory, extract point clouds and annotations by running `python collect_indoor3d_data.py`.
3. Enter the project root directory, generate training data by running
```bash
python tools/create_data.py s3dis --root-path ./data/s3dis --out-dir ./data/s3dis --extra-tag s3dis
```
The overall process could be achieved through the following script
```bash
python collect_indoor3d_data.py
cd ../..
python tools/create_data.py s3dis --root-path ./data/s3dis --out-dir ./data/s3dis --extra-tag s3dis
```
The directory structure after pre-processing should be as below
```
s3dis
├── indoor3d_util.py
├── collect_indoor3d_data.py
├── README.md
├── Stanford3dDataset_v1.2_Aligned_Version
├── s3dis_data
├── points
│ ├── xxxxx.bin
├── instance_mask
│ ├── xxxxx.bin
├── semantic_mask
│ ├── xxxxx.bin
├── seg_info
│ ├── Area_1_label_weight.npy
│ ├── Area_1_resampled_scene_idxs.npy
│ ├── Area_2_label_weight.npy
│ ├── Area_2_resampled_scene_idxs.npy
│ ├── Area_3_label_weight.npy
│ ├── Area_3_resampled_scene_idxs.npy
│ ├── Area_4_label_weight.npy
│ ├── Area_4_resampled_scene_idxs.npy
│ ├── Area_5_label_weight.npy
│ ├── Area_5_resampled_scene_idxs.npy
│ ├── Area_6_label_weight.npy
│ ├── Area_6_resampled_scene_idxs.npy
├── s3dis_infos_Area_1.pkl
├── s3dis_infos_Area_2.pkl
├── s3dis_infos_Area_3.pkl
├── s3dis_infos_Area_4.pkl
├── s3dis_infos_Area_5.pkl
├── s3dis_infos_Area_6.pkl
```
data/s3dis/collect_indoor3d_data.py 0 → 100644
View file @ a0090aa1
import argparse
import mmcv
from os import path as osp
from .indoor3d_util import export
parser = argparse.ArgumentParser()
parser.add_argument(
'--output-folder',
default='./s3dis_data',
help='output folder of the result.')
parser.add_argument(
'--data-dir',
default='Stanford3dDataset_v1.2_Aligned_Version',
help='s3dis data directory.')
parser.add_argument(
'--ann-file',
default='meta_data/anno_paths.txt',
help='The path of the file that stores the annotation names.')
args = parser.parse_args()
anno_paths = [line.rstrip() for line in open(args.ann_file)]
anno_paths = [osp.join(args.data_dir, p) for p in anno_paths]
output_folder = args.output_folder
mmcv.mkdir_or_exist(output_folder)
# Note: there is an extra character in the v1.2 data in Area_5/hallway_6.
# It's fixed manually here.
# Refer to https://github.com/AnTao97/dgcnn.pytorch/blob/843abe82dd731eb51a4b3f70632c2ed3c60560e9/prepare_data/collect_indoor3d_data.py#L18 # noqa
revise_file = osp.join(args.data_dir,
'Area_5/hallway_6/Annotations/ceiling_1.txt')
with open(revise_file, 'r') as f:
data = f.read()
# replace that extra character with blank space to separate data
data = data[:5545347] + ' ' + data[5545348:]
with open(revise_file, 'w') as f:
f.write(data)
for anno_path in anno_paths:
print(f'Exporting data from annotation file: {anno_path}')
elements = anno_path.split('/')
out_filename = \
elements[-3] + '_' + elements[-2] # Area_1_hallway_1
out_filename = osp.join(output_folder, out_filename)
if osp.isfile(f'{out_filename}_point.npy'):
print('File already exists. skipping.')
continue
export(anno_path, out_filename)
data/s3dis/indoor3d_util.py 0 → 100644
View file @ a0090aa1
import glob
import numpy as np
from os import path as osp
# -----------------------------------------------------------------------------
# CONSTANTS
# -----------------------------------------------------------------------------
BASE_DIR = osp.dirname(osp.abspath(__file__))
class_names = [
x.rstrip() for x in open(osp.join(BASE_DIR, 'meta_data/class_names.txt'))
]
class2label = {one_class: i for i, one_class in enumerate(class_names)}
# -----------------------------------------------------------------------------
# CONVERT ORIGINAL DATA TO POINTS, SEM_LABEL AND INS_LABEL FILES
# -----------------------------------------------------------------------------
def export(anno_path, out_filename):
"""Convert original dataset files to points, instance mask and semantic
mask files. We aggregated all the points from each instance in the room.
Args:
anno_path (str): path to annotations. e.g. Area_1/office_2/Annotations/
out_filename (str): path to save collected points and labels
file_format (str): txt or numpy, determines what file format to save.
Note:
the points are shifted before save, the most negative point is now
at origin.
"""
points_list = []
ins_idx = 1 # instance ids should be indexed from 1, so 0 is unannotated
for f in glob.glob(osp.join(anno_path, '*.txt')):
one_class = osp.basename(f).split('_')[0]
if one_class not in class_names: # some rooms have 'staris' class
one_class = 'clutter'
points = np.loadtxt(f)
labels = np.ones((points.shape[0], 1)) * class2label[one_class]
ins_labels = np.ones((points.shape[0], 1)) * ins_idx
ins_idx += 1
points_list.append(np.concatenate([points, labels, ins_labels], 1))
data_label = np.concatenate(points_list, 0) # [N, 8], (pts, rgb, sem, ins)
xyz_min = np.amin(data_label, axis=0)[0:3]
data_label[:, 0:3] -= xyz_min
np.save(f'{out_filename}_point.npy', data_label[:, :6].astype(np.float32))
np.save(f'{out_filename}_sem_label.npy', data_label[:, 6].astype(np.int))
np.save(f'{out_filename}_ins_label.npy', data_label[:, 7].astype(np.int))
data/s3dis/meta_data/anno_paths.txt 0 → 100644
View file @ a0090aa1
Area_1/conferenceRoom_1/Annotations
Area_1/conferenceRoom_2/Annotations
Area_1/copyRoom_1/Annotations
Area_1/hallway_1/Annotations
Area_1/hallway_2/Annotations
Area_1/hallway_3/Annotations
Area_1/hallway_4/Annotations
Area_1/hallway_5/Annotations
Area_1/hallway_6/Annotations
Area_1/hallway_7/Annotations
Area_1/hallway_8/Annotations
Area_1/office_10/Annotations
Area_1/office_11/Annotations
Area_1/office_12/Annotations
Area_1/office_13/Annotations
Area_1/office_14/Annotations
Area_1/office_15/Annotations
Area_1/office_16/Annotations
Area_1/office_17/Annotations
Area_1/office_18/Annotations
Area_1/office_19/Annotations
Area_1/office_1/Annotations
Area_1/office_20/Annotations
Area_1/office_21/Annotations
Area_1/office_22/Annotations
Area_1/office_23/Annotations
Area_1/office_24/Annotations
Area_1/office_25/Annotations
Area_1/office_26/Annotations
Area_1/office_27/Annotations
Area_1/office_28/Annotations
Area_1/office_29/Annotations
Area_1/office_2/Annotations
Area_1/office_30/Annotations
Area_1/office_31/Annotations
Area_1/office_3/Annotations
Area_1/office_4/Annotations
Area_1/office_5/Annotations
Area_1/office_6/Annotations
Area_1/office_7/Annotations
Area_1/office_8/Annotations
Area_1/office_9/Annotations
Area_1/pantry_1/Annotations
Area_1/WC_1/Annotations
Area_2/auditorium_1/Annotations
Area_2/auditorium_2/Annotations
Area_2/conferenceRoom_1/Annotations
Area_2/hallway_10/Annotations
Area_2/hallway_11/Annotations
Area_2/hallway_12/Annotations
Area_2/hallway_1/Annotations
Area_2/hallway_2/Annotations
Area_2/hallway_3/Annotations
Area_2/hallway_4/Annotations
Area_2/hallway_5/Annotations
Area_2/hallway_6/Annotations
Area_2/hallway_7/Annotations
Area_2/hallway_8/Annotations
Area_2/hallway_9/Annotations
Area_2/office_10/Annotations
Area_2/office_11/Annotations
Area_2/office_12/Annotations
Area_2/office_13/Annotations
Area_2/office_14/Annotations
Area_2/office_1/Annotations
Area_2/office_2/Annotations
Area_2/office_3/Annotations
Area_2/office_4/Annotations
Area_2/office_5/Annotations
Area_2/office_6/Annotations
Area_2/office_7/Annotations
Area_2/office_8/Annotations
Area_2/office_9/Annotations
Area_2/storage_1/Annotations
Area_2/storage_2/Annotations
Area_2/storage_3/Annotations
Area_2/storage_4/Annotations
Area_2/storage_5/Annotations
Area_2/storage_6/Annotations
Area_2/storage_7/Annotations
Area_2/storage_8/Annotations
Area_2/storage_9/Annotations
Area_2/WC_1/Annotations
Area_2/WC_2/Annotations
Area_3/conferenceRoom_1/Annotations
Area_3/hallway_1/Annotations
Area_3/hallway_2/Annotations
Area_3/hallway_3/Annotations
Area_3/hallway_4/Annotations
Area_3/hallway_5/Annotations
Area_3/hallway_6/Annotations
Area_3/lounge_1/Annotations
Area_3/lounge_2/Annotations
Area_3/office_10/Annotations
Area_3/office_1/Annotations
Area_3/office_2/Annotations
Area_3/office_3/Annotations
Area_3/office_4/Annotations
Area_3/office_5/Annotations
Area_3/office_6/Annotations
Area_3/office_7/Annotations
Area_3/office_8/Annotations
Area_3/office_9/Annotations
Area_3/storage_1/Annotations
Area_3/storage_2/Annotations
Area_3/WC_1/Annotations
Area_3/WC_2/Annotations
Area_4/conferenceRoom_1/Annotations
Area_4/conferenceRoom_2/Annotations
Area_4/conferenceRoom_3/Annotations
Area_4/hallway_10/Annotations
Area_4/hallway_11/Annotations
Area_4/hallway_12/Annotations
Area_4/hallway_13/Annotations
Area_4/hallway_14/Annotations
Area_4/hallway_1/Annotations
Area_4/hallway_2/Annotations
Area_4/hallway_3/Annotations
Area_4/hallway_4/Annotations
Area_4/hallway_5/Annotations
Area_4/hallway_6/Annotations
Area_4/hallway_7/Annotations
Area_4/hallway_8/Annotations
Area_4/hallway_9/Annotations
Area_4/lobby_1/Annotations
Area_4/lobby_2/Annotations
Area_4/office_10/Annotations
Area_4/office_11/Annotations
Area_4/office_12/Annotations
Area_4/office_13/Annotations
Area_4/office_14/Annotations
Area_4/office_15/Annotations
Area_4/office_16/Annotations
Area_4/office_17/Annotations
Area_4/office_18/Annotations
Area_4/office_19/Annotations
Area_4/office_1/Annotations
Area_4/office_20/Annotations
Area_4/office_21/Annotations
Area_4/office_22/Annotations
Area_4/office_2/Annotations
Area_4/office_3/Annotations
Area_4/office_4/Annotations
Area_4/office_5/Annotations
Area_4/office_6/Annotations
Area_4/office_7/Annotations
Area_4/office_8/Annotations
Area_4/office_9/Annotations
Area_4/storage_1/Annotations
Area_4/storage_2/Annotations
Area_4/storage_3/Annotations
Area_4/storage_4/Annotations
Area_4/WC_1/Annotations
Area_4/WC_2/Annotations
Area_4/WC_3/Annotations
Area_4/WC_4/Annotations
Area_5/conferenceRoom_1/Annotations
Area_5/conferenceRoom_2/Annotations
Area_5/conferenceRoom_3/Annotations
Area_5/hallway_10/Annotations
Area_5/hallway_11/Annotations
Area_5/hallway_12/Annotations
Area_5/hallway_13/Annotations
Area_5/hallway_14/Annotations
Area_5/hallway_15/Annotations
Area_5/hallway_1/Annotations
Area_5/hallway_2/Annotations
Area_5/hallway_3/Annotations
Area_5/hallway_4/Annotations
Area_5/hallway_5/Annotations
Area_5/hallway_6/Annotations
Area_5/hallway_7/Annotations
Area_5/hallway_8/Annotations
Area_5/hallway_9/Annotations
Area_5/lobby_1/Annotations
Area_5/office_10/Annotations
Area_5/office_11/Annotations
Area_5/office_12/Annotations
Area_5/office_13/Annotations
Area_5/office_14/Annotations
Area_5/office_15/Annotations
Area_5/office_16/Annotations
Area_5/office_17/Annotations
Area_5/office_18/Annotations
Area_5/office_19/Annotations
Area_5/office_1/Annotations
Area_5/office_20/Annotations
Area_5/office_21/Annotations
Area_5/office_22/Annotations
Area_5/office_23/Annotations
Area_5/office_24/Annotations
Area_5/office_25/Annotations
Area_5/office_26/Annotations
Area_5/office_27/Annotations
Area_5/office_28/Annotations
Area_5/office_29/Annotations
Area_5/office_2/Annotations
Area_5/office_30/Annotations
Area_5/office_31/Annotations
Area_5/office_32/Annotations
Area_5/office_33/Annotations
Area_5/office_34/Annotations
Area_5/office_35/Annotations
Area_5/office_36/Annotations
Area_5/office_37/Annotations
Area_5/office_38/Annotations
Area_5/office_39/Annotations
Area_5/office_3/Annotations
Area_5/office_40/Annotations
Area_5/office_41/Annotations
Area_5/office_42/Annotations
Area_5/office_4/Annotations
Area_5/office_5/Annotations
Area_5/office_6/Annotations
Area_5/office_7/Annotations
Area_5/office_8/Annotations
Area_5/office_9/Annotations
Area_5/pantry_1/Annotations
Area_5/storage_1/Annotations
Area_5/storage_2/Annotations
Area_5/storage_3/Annotations
Area_5/storage_4/Annotations
Area_5/WC_1/Annotations
Area_5/WC_2/Annotations
Area_6/conferenceRoom_1/Annotations
Area_6/copyRoom_1/Annotations
Area_6/hallway_1/Annotations
Area_6/hallway_2/Annotations
Area_6/hallway_3/Annotations
Area_6/hallway_4/Annotations
Area_6/hallway_5/Annotations
Area_6/hallway_6/Annotations
Area_6/lounge_1/Annotations
Area_6/office_10/Annotations
Area_6/office_11/Annotations
Area_6/office_12/Annotations
Area_6/office_13/Annotations
Area_6/office_14/Annotations
Area_6/office_15/Annotations
Area_6/office_16/Annotations
Area_6/office_17/Annotations
Area_6/office_18/Annotations
Area_6/office_19/Annotations
Area_6/office_1/Annotations
Area_6/office_20/Annotations
Area_6/office_21/Annotations
Area_6/office_22/Annotations
Area_6/office_23/Annotations
Area_6/office_24/Annotations
Area_6/office_25/Annotations
Area_6/office_26/Annotations
Area_6/office_27/Annotations
Area_6/office_28/Annotations
Area_6/office_29/Annotations
Area_6/office_2/Annotations
Area_6/office_30/Annotations
Area_6/office_31/Annotations
Area_6/office_32/Annotations
Area_6/office_33/Annotations
Area_6/office_34/Annotations
Area_6/office_35/Annotations
Area_6/office_36/Annotations
Area_6/office_37/Annotations
Area_6/office_3/Annotations
Area_6/office_4/Annotations
Area_6/office_5/Annotations
Area_6/office_6/Annotations
Area_6/office_7/Annotations
Area_6/office_8/Annotations
Area_6/office_9/Annotations
Area_6/openspace_1/Annotations
Area_6/pantry_1/Annotations
data/s3dis/meta_data/class_names.txt 0 → 100644
View file @ a0090aa1
ceiling
floor
wall
beam
column
window
door
table
chair
sofa
bookcase
board
clutter
docs/data_preparation.md
View file @ a0090aa1
...@@ -51,6 +51,12 @@ mmdetection3d ...@@ -51,6 +51,12 @@ mmdetection3d
│ │ ├── val.txt │ │ ├── val.txt
│ │ ├── test.txt │ │ ├── test.txt
│ │ ├── sample_submission.csv │ │ ├── sample_submission.csv
│ ├── s3dis
│ │ ├── meta_data
│ │ ├── Stanford3dDataset_v1.2_Aligned_Version
│ │ ├── collect_indoor3d_data.py
│ │ ├── indoor3d_util.py
│ │ ├── README.md
│ ├── scannet │ ├── scannet
│ │ ├── meta_data │ │ ├── meta_data
│ │ ├── scans │ │ ├── scans
...@@ -113,7 +119,9 @@ python tools/create_data.py lyft --root-path ./data/lyft --out-dir ./data/lyft - ...@@ -113,7 +119,9 @@ python tools/create_data.py lyft --root-path ./data/lyft --out-dir ./data/lyft -
Note that we follow the original folder names for clear organization. Please rename the raw folders as shown above. Note that we follow the original folder names for clear organization. Please rename the raw folders as shown above.
### ScanNet and SUN RGB-D ### S3DIS, ScanNet and SUN RGB-D
To prepare s3dis data, please see [s3dis](https://github.com/open-mmlab/mmdetection3d/blob/master/data/s3dis/README.md).
To prepare scannet data, please see [scannet](https://github.com/open-mmlab/mmdetection3d/blob/master/data/scannet/README.md). To prepare scannet data, please see [scannet](https://github.com/open-mmlab/mmdetection3d/blob/master/data/scannet/README.md).
......
mmdet3d/datasets/__init__.py
View file @ a0090aa1
...@@ -13,6 +13,7 @@ from .pipelines import (BackgroundPointsFilter, GlobalRotScaleTrans, ...@@ -13,6 +13,7 @@ from .pipelines import (BackgroundPointsFilter, GlobalRotScaleTrans,
NormalizePointsColor, ObjectNoise, ObjectRangeFilter, NormalizePointsColor, ObjectNoise, ObjectRangeFilter,
ObjectSample, PointShuffle, PointsRangeFilter, ObjectSample, PointShuffle, PointsRangeFilter,
RandomFlip3D, VoxelBasedPointSampler) RandomFlip3D, VoxelBasedPointSampler)
from .s3dis_dataset import S3DISSegDataset
from .scannet_dataset import ScanNetDataset, ScanNetSegDataset from .scannet_dataset import ScanNetDataset, ScanNetSegDataset
from .semantickitti_dataset import SemanticKITTIDataset from .semantickitti_dataset import SemanticKITTIDataset
from .sunrgbd_dataset import SUNRGBDDataset from .sunrgbd_dataset import SUNRGBDDataset
...@@ -27,7 +28,7 @@ __all__ = [ ...@@ -27,7 +28,7 @@ __all__ = [
'ObjectNoise', 'GlobalRotScaleTrans', 'PointShuffle', 'ObjectRangeFilter', 'ObjectNoise', 'GlobalRotScaleTrans', 'PointShuffle', 'ObjectRangeFilter',
'PointsRangeFilter', 'Collect3D', 'LoadPointsFromFile', 'PointsRangeFilter', 'Collect3D', 'LoadPointsFromFile',
'NormalizePointsColor', 'IndoorPointSample', 'LoadAnnotations3D', 'NormalizePointsColor', 'IndoorPointSample', 'LoadAnnotations3D',
'SUNRGBDDataset', 'ScanNetDataset', 'ScanNetSegDataset', 'SUNRGBDDataset', 'ScanNetDataset', 'ScanNetSegDataset', 'S3DISSegDataset',
'SemanticKITTIDataset', 'Custom3DDataset', 'Custom3DSegDataset', 'SemanticKITTIDataset', 'Custom3DDataset', 'Custom3DSegDataset',
'LoadPointsFromMultiSweeps', 'WaymoDataset', 'BackgroundPointsFilter', 'LoadPointsFromMultiSweeps', 'WaymoDataset', 'BackgroundPointsFilter',
'VoxelBasedPointSampler', 'get_loading_pipeline' 'VoxelBasedPointSampler', 'get_loading_pipeline'
......
mmdet3d/datasets/custom_3d_seg.py
View file @ a0090aa1
...@@ -268,6 +268,8 @@ class Custom3DSegDataset(Dataset): ...@@ -268,6 +268,8 @@ class Custom3DSegDataset(Dataset):
return np.arange(len(self.data_infos)).astype(np.int32), \ return np.arange(len(self.data_infos)).astype(np.int32), \
np.ones(len(self.CLASSES)).astype(np.float32) np.ones(len(self.CLASSES)).astype(np.float32)
# 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: if scene_idxs is None:
scene_idxs = np.arange(len(self.data_infos)) scene_idxs = np.arange(len(self.data_infos))
if isinstance(scene_idxs, str): if isinstance(scene_idxs, str):
...@@ -461,6 +463,10 @@ class Custom3DSegDataset(Dataset): ...@@ -461,6 +463,10 @@ class Custom3DSegDataset(Dataset):
def __getitem__(self, idx): def __getitem__(self, idx):
"""Get item from infos according to the given index. """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: Returns:
dict: Data dictionary of the corresponding index. dict: Data dictionary of the corresponding index.
""" """
......
mmdet3d/datasets/s3dis_dataset.py 0 → 100644
View file @ a0090aa1
import numpy as np
from os import path as osp
from mmdet3d.core import show_seg_result
from mmdet.datasets import DATASETS
from .custom_3d_seg import Custom3DSegDataset
from .pipelines import Compose
@DATASETS.register_module()
class _S3DISSegDataset(Custom3DSegDataset):
r"""S3DIS Dataset for Semantic Segmentation Task.
This class is the inner dataset for S3DIS. Since S3DIS has 6 areas, we
often train on 5 of them and test on the remaining one.
However, there is not a fixed train-test split of S3DIS. People often test
on Area_5 as suggested by `SEGCloud <https://arxiv.org/abs/1710.07563>`_.
But many papers also report the average results of 6-fold cross validation
over the 6 areas (e.g. `DGCNN <https://arxiv.org/abs/1801.07829>`_).
Therefore, we use an inner dataset for one area, and further use a dataset
wrapper to concat all the provided data in different areas.
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).
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.
label_weight (np.ndarray | str, optional): Precomputed weight to \
balance loss calculation. If None is given, compute from data.
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]]
def __init__(self,
data_root,
ann_file,
pipeline=None,
classes=None,
palette=None,
modality=None,
test_mode=False,
ignore_index=None,
scene_idxs=None,
label_weight=None):
super().__init__(
data_root=data_root,
ann_file=ann_file,
pipeline=pipeline,
classes=classes,
palette=palette,
modality=modality,
test_mode=test_mode,
ignore_index=ignore_index,
scene_idxs=scene_idxs,
label_weight=label_weight)
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),
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_and_label_weight(self, scene_idxs, label_weight):
"""Compute scene_idxs for data sampling and label weight for loss \
calculation.
We sample more times for scenes with more points. Label_weight is
inversely proportional to number of class points.
"""
# when testing, we load one whole scene every time
# and we don't need label weight for loss calculation
if not self.test_mode and scene_idxs is None:
raise NotImplementedError(
'please provide re-sampled scene indexes for training')
return super().get_scene_idxs_and_label_weight(scene_idxs,
label_weight)
@DATASETS.register_module()
class S3DISSegDataset(_S3DISSegDataset):
r"""S3DIS Dataset for Semantic Segmentation Task.
This class serves as the API for experiments on the S3DIS Dataset.
It wraps the provided datasets of different areas.
We don't use `mmdet.datasets.dataset_wrappers.ConcatDataset` because we
need to concat the `scene_idxs` and `label_weights` of different areas.
Please refer to the `google form <https://docs.google.com/forms/d/e/1FAIpQL
ScDimvNMCGhy_rmBA2gHfDu3naktRm6A8BPwAWWDv-Uhm6Shw/viewform?c=0&w=1>`_ for
data downloading.
Args:
data_root (str): Path of dataset root.
ann_files (list[str]): Path of several annotation files.
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).
Defaults to None.
scene_idxs (list[np.ndarray] | list[str], optional): Precomputed index
to load data. For scenes with many points, we may sample it several
times. Defaults to None.
label_weights (list[np.ndarray] | list[str], optional): Precomputed
weight to balance loss calculation. If None is given, compute from
data. Defaults to None.
"""
def __init__(self,
data_root,
ann_files,
pipeline=None,
classes=None,
palette=None,
modality=None,
test_mode=False,
ignore_index=None,
scene_idxs=None,
label_weights=None):
# make sure that ann_files, scene_idxs and label_weights have same len
ann_files = self._check_ann_files(ann_files)
scene_idxs = self._check_scene_idxs(scene_idxs, len(ann_files))
label_weights = self._check_label_weights(label_weights,
len(ann_files))
# initialize some attributes as datasets[0]
super().__init__(
data_root=data_root,
ann_file=ann_files[0],
pipeline=pipeline,
classes=classes,
palette=palette,
modality=modality,
test_mode=test_mode,
ignore_index=ignore_index,
scene_idxs=scene_idxs[0],
label_weight=label_weights[0])
datasets = [
_S3DISSegDataset(
data_root=data_root,
ann_file=ann_files[i],
pipeline=pipeline,
classes=classes,
palette=palette,
modality=modality,
test_mode=test_mode,
ignore_index=ignore_index,
scene_idxs=scene_idxs[i],
label_weight=label_weights[i]) for i in range(len(ann_files))
]
# data_infos, scene_idxs, label_weight need to be concat
self.concat_data_infos([dst.data_infos for dst in datasets])
self.concat_scene_idxs([dst.scene_idxs for dst in datasets])
self.concat_label_weight([dst.label_weight 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.
Args:
data_infos (list[list[dict]])
"""
self.data_infos = [
info for one_data_infos in data_infos for info in one_data_infos
]
def concat_scene_idxs(self, scene_idxs):
"""Concat scene_idxs from several datasets to form self.scene_idxs.
Needs to manually add offset to scene_idxs[1, 2, ...].
Args:
scene_idxs (list[np.ndarray])
"""
self.scene_idxs = np.array([], dtype=np.int32)
offset = 0
for one_scene_idxs in scene_idxs:
self.scene_idxs = np.concatenate(
[self.scene_idxs, one_scene_idxs + offset]).astype(np.int32)
offset = np.unique(self.scene_idxs).max() + 1
def concat_label_weight(self, label_weights):
"""Concat label_weight from several datasets to form self.label_weight.
Args:
label_weights (list[np.ndarray])
"""
# TODO: simply average them?
self.label_weight = np.array(label_weights).mean(0).astype(np.float32)
@staticmethod
def _duplicate_to_list(x, num):
"""Repeat x `num` times to form a list."""
return [x for _ in range(num)]
def _check_ann_files(self, ann_file):
"""Make ann_files as list/tuple."""
# ann_file could be str
if not isinstance(ann_file, (list, tuple)):
ann_file = self._duplicate_to_list(ann_file, 1)
return ann_file
def _check_scene_idxs(self, scene_idx, num):
"""Make scene_idxs as list/tuple."""
if scene_idx is None:
return self._duplicate_to_list(scene_idx, num)
# scene_idx could be str, np.ndarray, list or tuple
if isinstance(scene_idx, str): # str
return self._duplicate_to_list(scene_idx, num)
if isinstance(scene_idx[0], str): # list of str
return scene_idx
if isinstance(scene_idx[0], (list, tuple, np.ndarray)): # list of idx
return scene_idx
# single idx
return self._duplicate_to_list(scene_idx, num)
def _check_label_weights(self, label_weight, num):
"""Make label_weights as list/tuple."""
if label_weight is None:
return self._duplicate_to_list(label_weight, num)
# label_weight could be str, np.ndarray, list or tuple
if isinstance(label_weight, str): # str
return self._duplicate_to_list(label_weight, num)
if isinstance(label_weight[0], str): # list of str
return label_weight
if isinstance(label_weight[0], (list, tuple, np.ndarray)): # list of w
return label_weight
# single weight
return self._duplicate_to_list(label_weight, num)
mmdet3d/datasets/scannet_dataset.py
View file @ a0090aa1
...@@ -276,8 +276,7 @@ class ScanNetSegDataset(Custom3DSegDataset): ...@@ -276,8 +276,7 @@ class ScanNetSegDataset(Custom3DSegDataset):
with_seg_3d=True), with_seg_3d=True),
dict( dict(
type='PointSegClassMapping', type='PointSegClassMapping',
valid_cat_ids=(1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, valid_cat_ids=self.VALID_CLASS_IDS),
24, 28, 33, 34, 36, 39)),
dict( dict(
type='DefaultFormatBundle3D', type='DefaultFormatBundle3D',
with_label=False, with_label=False,
......
tests/test_data/test_datasets/test_s3dis_dataset.py 0 → 100644
View file @ a0090aa1
import numpy as np
import pytest
import torch
from mmdet3d.datasets import S3DISSegDataset
def test_seg_getitem():
np.random.seed(0)
root_path = './tests/data/s3dis/'
ann_file = './tests/data/s3dis/s3dis_infos.pkl'
class_names = ('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)]
label_weight = [
3.0441623, 3.3606708, 2.6408234, 4.5086737, 4.8403897, 4.7637715,
4.4685297, 4.7051463, 4.9190116, 5.3899403, 4.6436925, 5.0669650,
3.6270046
]
pipelines = [
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=tuple(range(len(class_names)))),
dict(
type='IndoorPatchPointSample',
num_points=5,
block_size=1.0,
sample_rate=1.0,
ignore_index=len(class_names),
use_normalized_coord=True),
dict(type='NormalizePointsColor', color_mean=None),
dict(type='DefaultFormatBundle3D', class_names=class_names),
dict(
type='Collect3D',
keys=['points', 'pts_semantic_mask'],
meta_keys=['file_name', 'sample_idx'])
]
s3dis_dataset = S3DISSegDataset(
data_root=root_path,
ann_files=ann_file,
pipeline=pipelines,
classes=None,
palette=None,
modality=None,
test_mode=False,
ignore_index=None,
scene_idxs=scene_idxs,
label_weights=label_weight)
data = s3dis_dataset[0]
points = data['points']._data
pts_semantic_mask = data['pts_semantic_mask']._data
file_name = data['img_metas']._data['file_name']
sample_idx = data['img_metas']._data['sample_idx']
assert file_name == './tests/data/s3dis/points/Area_1_office_2.bin'
assert sample_idx == 'Area_1_office_2'
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])
original_classes = s3dis_dataset.CLASSES
original_palette = s3dis_dataset.PALETTE
assert s3dis_dataset.CLASSES == class_names
assert s3dis_dataset.ignore_index == 13
assert torch.allclose(points, expected_points, 1e-2)
assert np.all(pts_semantic_mask.numpy() == expected_pts_semantic_mask)
assert original_classes == class_names
assert original_palette == palette
assert s3dis_dataset.scene_idxs.dtype == np.int32
assert np.all(s3dis_dataset.scene_idxs == np.array(scene_idxs))
assert np.allclose(s3dis_dataset.label_weight, np.array(label_weight),
1e-5)
# test dataset with selected classes
s3dis_dataset = S3DISSegDataset(
data_root=root_path,
ann_files=ann_file,
pipeline=None,
classes=['beam', 'window'],
scene_idxs=scene_idxs)
label_map = {i: 13 for i in range(14)}
label_map.update({3: 0, 5: 1})
assert s3dis_dataset.CLASSES != original_classes
assert s3dis_dataset.CLASSES == ['beam', 'window']
assert s3dis_dataset.PALETTE == [palette[3], palette[5]]
assert s3dis_dataset.VALID_CLASS_IDS == [3, 5]
assert s3dis_dataset.label_map == label_map
assert s3dis_dataset.label2cat == {0: 'beam', 1: 'window'}
assert np.all(s3dis_dataset.label_weight == np.ones(2))
# test load classes from file
import tempfile
tmp_file = tempfile.NamedTemporaryFile()
with open(tmp_file.name, 'w') as f:
f.write('beam\nwindow\n')
s3dis_dataset = S3DISSegDataset(
data_root=root_path,
ann_files=ann_file,
pipeline=None,
classes=tmp_file.name,
scene_idxs=scene_idxs)
assert s3dis_dataset.CLASSES != original_classes
assert s3dis_dataset.CLASSES == ['beam', 'window']
assert s3dis_dataset.PALETTE == [palette[3], palette[5]]
assert s3dis_dataset.VALID_CLASS_IDS == [3, 5]
assert s3dis_dataset.label_map == label_map
assert s3dis_dataset.label2cat == {0: 'beam', 1: 'window'}
# test scene_idxs in dataset
# we should input scene_idxs in train mode
with pytest.raises(NotImplementedError):
s3dis_dataset = S3DISSegDataset(
data_root=root_path,
ann_files=ann_file,
pipeline=None,
scene_idxs=None)
# test mode
s3dis_dataset = S3DISSegDataset(
data_root=root_path,
ann_files=ann_file,
pipeline=None,
test_mode=True,
scene_idxs=scene_idxs)
assert np.all(s3dis_dataset.scene_idxs == np.array([0]))
assert np.all(s3dis_dataset.label_weight == np.ones(len(class_names)))
def test_seg_evaluate():
if not torch.cuda.is_available():
pytest.skip()
root_path = './tests/data/s3dis'
ann_file = './tests/data/s3dis/s3dis_infos.pkl'
s3dis_dataset = S3DISSegDataset(
data_root=root_path, ann_files=ann_file, test_mode=True)
results = []
pred_sem_mask = dict(
semantic_mask=torch.tensor([
2, 3, 1, 2, 2, 6, 1, 0, 1, 1, 9, 12, 3, 0, 2, 0, 2, 0, 8, 3, 1, 2,
0, 2, 1, 7, 2, 10, 2, 0, 0, 0, 2, 3, 2, 2, 2, 2, 2, 3, 0, 0, 4, 6,
7, 2, 1, 2, 0, 1, 7, 0, 2, 2, 2, 0, 2, 2, 1, 12, 0, 2, 2, 2, 2, 7,
2, 2, 0, 2, 6, 2, 12, 6, 3, 12, 2, 1, 6, 1, 2, 6, 8, 2, 10, 1, 11,
0, 6, 9, 4, 3, 0, 0, 12, 1, 1, 5, 3, 2
]).long())
results.append(pred_sem_mask)
ret_dict = s3dis_dataset.evaluate(results)
assert abs(ret_dict['miou'] - 0.7625) < 0.01
assert abs(ret_dict['acc'] - 0.9) < 0.01
assert abs(ret_dict['acc_cls'] - 0.9074) < 0.01
def test_seg_show():
import mmcv
import tempfile
from os import path as osp
tmp_dir = tempfile.TemporaryDirectory()
temp_dir = tmp_dir.name
root_path = './tests/data/s3dis'
ann_file = './tests/data/s3dis/s3dis_infos.pkl'
s3dis_dataset = S3DISSegDataset(
data_root=root_path, ann_files=ann_file, scene_idxs=[0])
result = dict(
semantic_mask=torch.tensor([
2, 2, 1, 2, 2, 5, 1, 0, 1, 1, 9, 12, 3, 0, 2, 0, 2, 0, 8, 2, 0, 2,
0, 2, 1, 7, 2, 10, 2, 0, 0, 0, 2, 2, 2, 2, 2, 1, 2, 2, 0, 0, 4, 6,
7, 2, 1, 2, 0, 1, 7, 0, 2, 2, 2, 0, 2, 2, 1, 12, 0, 2, 2, 2, 2, 7,
2, 2, 0, 2, 6, 2, 12, 6, 2, 12, 2, 1, 6, 1, 2, 6, 8, 2, 10, 1, 10,
0, 6, 9, 4, 3, 0, 0, 12, 1, 1, 5, 2, 2
]).long())
results = [result]
s3dis_dataset.show(results, temp_dir, show=False)
pts_file_path = osp.join(temp_dir, 'Area_1_office_2',
'Area_1_office_2_points.obj')
gt_file_path = osp.join(temp_dir, 'Area_1_office_2',
'Area_1_office_2_gt.obj')
pred_file_path = osp.join(temp_dir, 'Area_1_office_2',
'Area_1_office_2_pred.obj')
mmcv.check_file_exist(pts_file_path)
mmcv.check_file_exist(gt_file_path)
mmcv.check_file_exist(pred_file_path)
tmp_dir.cleanup()
def test_multi_areas():
# S3DIS dataset has 6 areas, we often train on several of them
# need to verify the concat function of S3DISSegDataset
root_path = './tests/data/s3dis'
ann_file = './tests/data/s3dis/s3dis_infos.pkl'
class_names = ('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)]
label_weight = [
3.0441623, 3.3606708, 2.6408234, 4.5086737, 4.8403897, 4.7637715,
4.4685297, 4.7051463, 4.9190116, 5.3899403, 4.6436925, 5.0669650,
3.6270046
]
# repeat
repeat_num = 3
s3dis_dataset = S3DISSegDataset(
data_root=root_path,
ann_files=[ann_file for _ in range(repeat_num)],
scene_idxs=scene_idxs,
label_weights=label_weight)
assert s3dis_dataset.CLASSES == class_names
assert s3dis_dataset.PALETTE == palette
assert len(s3dis_dataset.data_infos) == repeat_num
assert np.all(s3dis_dataset.scene_idxs == np.concatenate(
[np.array(scene_idxs) + i for i in range(repeat_num)]))
assert np.allclose(s3dis_dataset.label_weight, np.array(label_weight))
# different scene_idxs and label_weight input
label_weights = np.random.rand(repeat_num, len(class_names))
s3dis_dataset = S3DISSegDataset(
data_root=root_path,
ann_files=[ann_file for _ in range(repeat_num)],
scene_idxs=[[0, 0, 1, 2, 2], [0, 1, 2, 3, 3, 4], [0, 1, 1, 2, 2, 2]],
label_weights=label_weights)
assert np.all(s3dis_dataset.scene_idxs == np.array(
[0, 0, 1, 2, 2, 3, 4, 5, 6, 6, 7, 8, 9, 9, 10, 10, 10]))
assert np.allclose(s3dis_dataset.label_weight, label_weights.mean(0))
tests/test_data/test_pipelines/test_indoor_pipeline.py
View file @ a0090aa1
...@@ -175,6 +175,75 @@ def test_scannet_seg_pipeline(): ...@@ -175,6 +175,75 @@ def test_scannet_seg_pipeline():
assert np.all(pts_semantic_mask.numpy() == expected_pts_semantic_mask) assert np.all(pts_semantic_mask.numpy() == expected_pts_semantic_mask)
def test_s3dis_seg_pipeline():
class_names = ('ceiling', 'floor', 'wall', 'beam', 'column', 'window',
'door', 'table', 'chair', 'sofa', 'bookcase', 'board',
'clutter')
np.random.seed(0)
pipelines = [
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=tuple(range(len(class_names)))),
dict(
type='IndoorPatchPointSample',
num_points=5,
block_size=1.0,
sample_rate=1.0,
ignore_index=len(class_names),
use_normalized_coord=True),
dict(type='NormalizePointsColor', color_mean=None),
dict(type='DefaultFormatBundle3D', class_names=class_names),
dict(type='Collect3D', keys=['points', 'pts_semantic_mask'])
]
pipeline = Compose(pipelines)
info = mmcv.load('./tests/data/s3dis/s3dis_infos.pkl')[0]
results = dict()
data_path = './tests/data/s3dis'
results['pts_filename'] = osp.join(data_path, info['pts_path'])
results['ann_info'] = dict()
results['ann_info']['pts_semantic_mask_path'] = osp.join(
data_path, info['pts_semantic_mask_path'])
results['pts_seg_fields'] = []
results = pipeline(results)
points = results['points']._data
pts_semantic_mask = results['pts_semantic_mask']._data
# build sampled points
s3dis_points = np.fromfile(
osp.join(data_path, info['pts_path']), dtype=np.float32).reshape(
(-1, 6))
s3dis_choices = np.array([87, 37, 60, 18, 31])
s3dis_center = np.array([2.691, 2.231, 3.172])
s3dis_center[2] = 0.0
s3dis_coord_max = np.amax(s3dis_points[:, :3], axis=0)
expected_points = np.concatenate([
s3dis_points[s3dis_choices, :3] - s3dis_center,
s3dis_points[s3dis_choices, 3:] / 255.,
s3dis_points[s3dis_choices, :3] / s3dis_coord_max
],
axis=1)
expected_pts_semantic_mask = np.array([0, 1, 0, 8, 0])
assert np.allclose(points.numpy(), expected_points, atol=1e-6)
assert np.all(pts_semantic_mask.numpy() == expected_pts_semantic_mask)
def test_sunrgbd_pipeline(): def test_sunrgbd_pipeline():
class_names = ('bed', 'table', 'sofa', 'chair', 'toilet', 'desk', class_names = ('bed', 'table', 'sofa', 'chair', 'toilet', 'desk',
'dresser', 'night_stand', 'bookshelf', 'bathtub') 'dresser', 'night_stand', 'bookshelf', 'bathtub')
......
tests/test_data/test_pipelines/test_indoor_sample.py
View file @ a0090aa1
...@@ -96,8 +96,7 @@ def test_indoor_seg_sample(): ...@@ -96,8 +96,7 @@ def test_indoor_seg_sample():
scannet_points[scannet_choices, :3] - scannet_center, scannet_points[scannet_choices, :3] - scannet_center,
scannet_points[scannet_choices, 3:], scannet_points[scannet_choices, 3:],
scannet_points[scannet_choices, :3] / scannet_coord_max scannet_points[scannet_choices, :3] / scannet_coord_max
], ], 1)
axis=1)
assert scannet_points_result.points_dim == 9 assert scannet_points_result.points_dim == 9
assert scannet_points_result.attribute_dims == dict( assert scannet_points_result.attribute_dims == dict(
...@@ -115,3 +114,39 @@ def test_indoor_seg_sample(): ...@@ -115,3 +114,39 @@ def test_indoor_seg_sample():
'use_normalized_coord=True, ' \ 'use_normalized_coord=True, ' \
'num_try=10)' 'num_try=10)'
assert repr_str == expected_repr_str assert repr_str == expected_repr_str
# test on S3DIS dataset
np.random.seed(0)
s3dis_patch_sample_points = IndoorPatchPointSample(5, 1.0, 1.0, None, True)
s3dis_results = dict()
s3dis_points = np.fromfile(
'./tests/data/s3dis/points/Area_1_office_2.bin',
dtype=np.float32).reshape((-1, 6))
s3dis_results['points'] = DepthPoints(
s3dis_points, points_dim=6, attribute_dims=dict(color=[3, 4, 5]))
s3dis_pts_semantic_mask = np.fromfile(
'./tests/data/s3dis/semantic_mask/Area_1_office_2.bin', dtype=np.long)
s3dis_results['pts_semantic_mask'] = s3dis_pts_semantic_mask
s3dis_results = s3dis_patch_sample_points(s3dis_results)
s3dis_points_result = s3dis_results['points']
s3dis_semantic_labels_result = s3dis_results['pts_semantic_mask']
# manually constructed sampled points
s3dis_choices = np.array([87, 37, 60, 18, 31])
s3dis_center = np.array([2.691, 2.231, 3.172])
s3dis_center[2] = 0.0
s3dis_coord_max = np.amax(s3dis_points[:, :3], axis=0)
s3dis_input_points = np.concatenate([
s3dis_points[s3dis_choices, :3] - s3dis_center,
s3dis_points[s3dis_choices,
3:], s3dis_points[s3dis_choices, :3] / s3dis_coord_max
], 1)
assert s3dis_points_result.points_dim == 9
assert s3dis_points_result.attribute_dims == dict(
color=[3, 4, 5], normalized_coord=[6, 7, 8])
s3dis_points_result = s3dis_points_result.tensor.numpy()
assert np.allclose(s3dis_input_points, s3dis_points_result, atol=1e-6)
assert np.all(np.array([0, 1, 0, 8, 0]) == s3dis_semantic_labels_result)
tests/test_data/test_pipelines/test_loadings/test_loading.py
View file @ a0090aa1
...@@ -12,6 +12,7 @@ from mmdet3d.datasets.pipelines import (LoadAnnotations3D, LoadPointsFromFile, ...@@ -12,6 +12,7 @@ from mmdet3d.datasets.pipelines import (LoadAnnotations3D, LoadPointsFromFile,
def test_load_points_from_indoor_file(): def test_load_points_from_indoor_file():
# test on SUN RGB-D dataset with shifted height
sunrgbd_info = mmcv.load('./tests/data/sunrgbd/sunrgbd_infos.pkl') sunrgbd_info = mmcv.load('./tests/data/sunrgbd/sunrgbd_infos.pkl')
sunrgbd_load_points_from_file = LoadPointsFromFile( sunrgbd_load_points_from_file = LoadPointsFromFile(
coord_type='DEPTH', load_dim=6, shift_height=True) coord_type='DEPTH', load_dim=6, shift_height=True)
...@@ -31,6 +32,7 @@ def test_load_points_from_indoor_file(): ...@@ -31,6 +32,7 @@ def test_load_points_from_indoor_file():
data_path = './tests/data/scannet' data_path = './tests/data/scannet'
scannet_info = scannet_info[0] scannet_info = scannet_info[0]
# test on ScanNet dataset with shifted height
scannet_results['pts_filename'] = osp.join(data_path, scannet_results['pts_filename'] = osp.join(data_path,
scannet_info['pts_path']) scannet_info['pts_path'])
scannet_results = scannet_load_data(scannet_results) scannet_results = scannet_load_data(scannet_results)
...@@ -64,6 +66,28 @@ def test_load_points_from_indoor_file(): ...@@ -64,6 +66,28 @@ def test_load_points_from_indoor_file():
scannet_point_cloud = scannet_point_cloud.tensor.numpy() scannet_point_cloud = scannet_point_cloud.tensor.numpy()
assert scannet_point_cloud.shape == (100, 7) assert scannet_point_cloud.shape == (100, 7)
# test load point cloud on S3DIS with color
data_path = './tests/data/s3dis'
s3dis_info = mmcv.load('./tests/data/s3dis/s3dis_infos.pkl')
s3dis_info = s3dis_info[0]
s3dis_load_data = LoadPointsFromFile(
coord_type='DEPTH',
load_dim=6,
use_dim=[0, 1, 2, 3, 4, 5],
shift_height=False,
use_color=True)
s3dis_results = dict()
s3dis_results['pts_filename'] = osp.join(data_path, s3dis_info['pts_path'])
s3dis_results = s3dis_load_data(s3dis_results)
s3dis_point_cloud = s3dis_results['points']
assert s3dis_point_cloud.points_dim == 6
assert s3dis_point_cloud.attribute_dims == dict(color=[3, 4, 5])
s3dis_point_cloud = s3dis_point_cloud.tensor.numpy()
assert s3dis_point_cloud.shape == (100, 6)
def test_load_points_from_outdoor_file(): def test_load_points_from_outdoor_file():
data_path = 'tests/data/kitti/a.bin' data_path = 'tests/data/kitti/a.bin'
...@@ -141,6 +165,33 @@ def test_load_annotations3D(): ...@@ -141,6 +165,33 @@ def test_load_annotations3D():
assert scannet_pts_instance_mask.shape == (100, ) assert scannet_pts_instance_mask.shape == (100, )
assert scannet_pts_semantic_mask.shape == (100, ) assert scannet_pts_semantic_mask.shape == (100, )
# Test s3dis LoadAnnotations3D
s3dis_info = mmcv.load('./tests/data/s3dis/s3dis_infos.pkl')[0]
s3dis_load_annotations3D = LoadAnnotations3D(
with_bbox_3d=False,
with_label_3d=False,
with_mask_3d=True,
with_seg_3d=True)
s3dis_results = dict()
data_path = './tests/data/s3dis'
# prepare input of loading pipeline
s3dis_results['ann_info'] = dict()
s3dis_results['ann_info']['pts_instance_mask_path'] = osp.join(
data_path, s3dis_info['pts_instance_mask_path'])
s3dis_results['ann_info']['pts_semantic_mask_path'] = osp.join(
data_path, s3dis_info['pts_semantic_mask_path'])
s3dis_results['pts_mask_fields'] = []
s3dis_results['pts_seg_fields'] = []
s3dis_results = s3dis_load_annotations3D(s3dis_results)
s3dis_pts_instance_mask = s3dis_results['pts_instance_mask']
s3dis_pts_semantic_mask = s3dis_results['pts_semantic_mask']
assert s3dis_pts_instance_mask.shape == (100, )
assert s3dis_pts_semantic_mask.shape == (100, )
def test_load_segmentation_mask(): def test_load_segmentation_mask():
# Test loading semantic segmentation mask on ScanNet dataset # Test loading semantic segmentation mask on ScanNet dataset
...@@ -178,6 +229,39 @@ def test_load_segmentation_mask(): ...@@ -178,6 +229,39 @@ def test_load_segmentation_mask():
12, 2, 20, 0, 0, 13, 20, 1, 20, 5, 3, 0, 13, 1, 2, 2, 2, 1 12, 2, 20, 0, 0, 13, 20, 1, 20, 5, 3, 0, 13, 1, 2, 2, 2, 1
])) ]))
# Test on S3DIS dataset
s3dis_info = mmcv.load('./tests/data/s3dis/s3dis_infos.pkl')[0]
s3dis_load_annotations3D = LoadAnnotations3D(
with_bbox_3d=False,
with_label_3d=False,
with_mask_3d=False,
with_seg_3d=True)
s3dis_results = dict()
data_path = './tests/data/s3dis'
# prepare input of loading pipeline
s3dis_results['ann_info'] = dict()
s3dis_results['ann_info']['pts_semantic_mask_path'] = osp.join(
data_path, s3dis_info['pts_semantic_mask_path'])
s3dis_results['pts_seg_fields'] = []
s3dis_results = s3dis_load_annotations3D(s3dis_results)
s3dis_pts_semantic_mask = s3dis_results['pts_semantic_mask']
assert s3dis_pts_semantic_mask.shape == (100, )
# Convert class_id to label and assign ignore_index
s3dis_seg_class_mapping = PointSegClassMapping(tuple(range(13)))
s3dis_results = s3dis_seg_class_mapping(s3dis_results)
s3dis_pts_semantic_mask = s3dis_results['pts_semantic_mask']
assert np.all(s3dis_pts_semantic_mask == np.array([
2, 2, 1, 2, 2, 5, 1, 0, 1, 1, 9, 12, 3, 0, 2, 0, 2, 0, 8, 2, 0, 2, 0,
2, 1, 7, 2, 10, 2, 0, 0, 0, 2, 2, 2, 2, 2, 1, 2, 2, 0, 0, 4, 6, 7, 2,
1, 2, 0, 1, 7, 0, 2, 2, 2, 0, 2, 2, 1, 12, 0, 2, 2, 2, 2, 7, 2, 2, 0,
2, 6, 2, 12, 6, 2, 12, 2, 1, 6, 1, 2, 6, 8, 2, 10, 1, 10, 0, 6, 9, 4,
3, 0, 0, 12, 1, 1, 5, 2, 2
]))
def test_load_points_from_multi_sweeps(): def test_load_points_from_multi_sweeps():
load_points_from_multi_sweeps = LoadPointsFromMultiSweeps() load_points_from_multi_sweeps = LoadPointsFromMultiSweeps()
......
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