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/
# demo
*.jpg
*.png
/data/s3dis/Stanford3dDataset_v1.2_Aligned_Version/
/data/scannet/scans/
/data/sunrgbd/OFFICIAL_SUNRGBD/
*.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
│ │ ├── val.txt
│ │ ├── test.txt
│ │ ├── sample_submission.csv
│ ├── s3dis
│ │ ├── meta_data
│ │ ├── Stanford3dDataset_v1.2_Aligned_Version
│ │ ├── collect_indoor3d_data.py
│ │ ├── indoor3d_util.py
│ │ ├── README.md
│ ├── scannet
│ │ ├── meta_data
│ │ ├── scans
......@@ -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.
### 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).
......
mmdet3d/datasets/__init__.py
View file @ a0090aa1
......@@ -13,6 +13,7 @@ from .pipelines import (BackgroundPointsFilter, GlobalRotScaleTrans,
NormalizePointsColor, ObjectNoise, ObjectRangeFilter,
ObjectSample, PointShuffle, PointsRangeFilter,
RandomFlip3D, VoxelBasedPointSampler)
from .s3dis_dataset import S3DISSegDataset
from .scannet_dataset import ScanNetDataset, ScanNetSegDataset
from .semantickitti_dataset import SemanticKITTIDataset
from .sunrgbd_dataset import SUNRGBDDataset
......@@ -27,7 +28,7 @@ __all__ = [
'ObjectNoise', 'GlobalRotScaleTrans', 'PointShuffle', 'ObjectRangeFilter',
'PointsRangeFilter', 'Collect3D', 'LoadPointsFromFile',
'NormalizePointsColor', 'IndoorPointSample', 'LoadAnnotations3D',
'SUNRGBDDataset', 'ScanNetDataset', 'ScanNetSegDataset',
'SUNRGBDDataset', 'ScanNetDataset', 'ScanNetSegDataset', 'S3DISSegDataset',
'SemanticKITTIDataset', 'Custom3DDataset', 'Custom3DSegDataset',
'LoadPointsFromMultiSweeps', 'WaymoDataset', 'BackgroundPointsFilter',
'VoxelBasedPointSampler', 'get_loading_pipeline'
......
mmdet3d/datasets/custom_3d_seg.py
View file @ a0090aa1
......@@ -268,6 +268,8 @@ class Custom3DSegDataset(Dataset):
return np.arange(len(self.data_infos)).astype(np.int32), \
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:
scene_idxs = np.arange(len(self.data_infos))
if isinstance(scene_idxs, str):
......@@ -461,6 +463,10 @@ class Custom3DSegDataset(Dataset):
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.
"""
......
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):
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)),
valid_cat_ids=self.VALID_CLASS_IDS),
dict(
type='DefaultFormatBundle3D',
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():
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():
class_names = ('bed', 'table', 'sofa', 'chair', 'toilet', 'desk',
'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():
scannet_points[scannet_choices, :3] - scannet_center,
scannet_points[scannet_choices, 3:],
scannet_points[scannet_choices, :3] / scannet_coord_max
],
axis=1)
], 1)
assert scannet_points_result.points_dim == 9
assert scannet_points_result.attribute_dims == dict(
......@@ -115,3 +114,39 @@ def test_indoor_seg_sample():
'use_normalized_coord=True, ' \
'num_try=10)'
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,
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_load_points_from_file = LoadPointsFromFile(
coord_type='DEPTH', load_dim=6, shift_height=True)
......@@ -31,6 +32,7 @@ def test_load_points_from_indoor_file():
data_path = './tests/data/scannet'
scannet_info = scannet_info[0]
# test on ScanNet dataset with shifted height
scannet_results['pts_filename'] = osp.join(data_path,
scannet_info['pts_path'])
scannet_results = scannet_load_data(scannet_results)
......@@ -64,6 +66,28 @@ def test_load_points_from_indoor_file():
scannet_point_cloud = scannet_point_cloud.tensor.numpy()
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():
data_path = 'tests/data/kitti/a.bin'
......@@ -141,6 +165,33 @@ def test_load_annotations3D():
assert scannet_pts_instance_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():
# Test loading semantic segmentation mask on ScanNet dataset
......@@ -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
]))
# 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():
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