[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

.gitignore

@@ -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,0 +1,116 @@
+# 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,0 +1,58 @@
+### 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,0 +1,49 @@
+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,0 +1,53 @@
+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))