Merge branch 'indoor_data_converter' into 'master'

Indoor data converter

See merge request open-mmlab/mmdet.3d!2

.isort.cfg

@@ -3,6 +3,6 @@ line_length = 79
 multi_line_output = 0
 known_standard_library = setuptools
 known_first_party = mmdet,mmdet3d
-known_third_party = cv2,mmcv,numba,numpy,nuscenes,pycocotools,pyquaternion,pytest,shapely,six,skimage,torch,torchvision
+known_third_party = cv2,mmcv,numba,numpy,nuscenes,pycocotools,pyquaternion,pytest,scipy,shapely,six,skimage,torch,torchvision
 no_lines_before = STDLIB,LOCALFOLDER
 default_section = THIRDPARTY

tools/create_data.py

@@ -3,6 +3,8 @@ import os.path as osp
 
 import tools.data_converter.kitti_converter as kitti
 import tools.data_converter.nuscenes_converter as nuscenes_converter
+import tools.data_converter.scannet_converter as scannet
+import tools.data_converter.sunrgbd_converter as sunrgbd
 from tools.data_converter.create_gt_database import create_groundtruth_database
 
 
@@ -43,6 +45,14 @@ def nuscenes_data_prep(root_path,
         '{}/{}_infos_train.pkl'.format(out_dir, info_prefix))
 
 
+def scannet_data_prep(root_path, info_prefix, out_dir):
+    scannet.create_scannet_info_file(root_path, info_prefix, out_dir)
+
+
+def sunrgbd_data_prep(root_path, info_prefix, out_dir):
+    sunrgbd.create_sunrgbd_info_file(root_path, info_prefix, out_dir)
+
+
 parser = argparse.ArgumentParser(description='Data converter arg parser')
 parser.add_argument('dataset', metavar='kitti', help='name of the dataset')
 parser.add_argument(
@@ -104,3 +114,13 @@ if __name__ == '__main__':
             dataset_name='NuScenesDataset',
             out_dir=args.out_dir,
             max_sweeps=args.max_sweeps)
+    elif args.dataset == 'scannet':
+        scannet_data_prep(
+            root_path=args.root_path,
+            info_prefix=args.extra_tag,
+            out_dir=args.out_dir)
+    elif args.dataset == 'sunrgbd':
+        sunrgbd_data_prep(
+            root_path=args.root_path,
+            info_prefix=args.extra_tag,
+            out_dir=args.out_dir)

tools/data_converter/scannet_converter.py

+import os
+import pickle
+from pathlib import Path
+
+from tools.data_converter.scannet_data_utils import ScanNetData
+
+
+def create_scannet_info_file(data_path, pkl_prefix='scannet', save_path=None):
+    assert os.path.exists(data_path)
+    if save_path is None:
+        save_path = Path(data_path)
+    else:
+        save_path = Path(save_path)
+    assert os.path.exists(save_path)
+    train_filename = save_path / f'{pkl_prefix}_infos_train.pkl'
+    val_filename = save_path / f'{pkl_prefix}_infos_val.pkl'
+    train_dataset = ScanNetData(root_path=data_path, split='train')
+    val_dataset = ScanNetData(root_path=data_path, split='val')
+    scannet_infos_train = train_dataset.get_scannet_infos(has_label=True)
+    with open(train_filename, 'wb') as f:
+        pickle.dump(scannet_infos_train, f)
+    print('Scannet info train file is saved to %s' % train_filename)
+    scannet_infos_val = val_dataset.get_scannet_infos(has_label=True)
+    with open(val_filename, 'wb') as f:
+        pickle.dump(scannet_infos_val, f)
+    print('Scannet info val file is saved to %s' % val_filename)
+
+
+if __name__ == '__main__':
+    create_scannet_info_file(
+        data_path='./data/scannet', save_path='./data/scannet')

tools/data_converter/scannet_data_utils.py

+import os
+
+import numpy as np
+
+
+class ScanNetData(object):
+    ''' Load and parse object data '''
+
+    def __init__(self, root_path, split='train'):
+        self.root_dir = root_path
+        self.split = split
+        self.split_dir = os.path.join(root_path)
+        self.classes = [
+            'cabinet', 'bed', 'chair', 'sofa', 'table', 'door', 'window',
+            'bookshelf', 'picture', 'counter', 'desk', 'curtain',
+            'refrigerator', 'showercurtrain', 'toilet', 'sink', 'bathtub',
+            'garbagebin'
+        ]
+        self.cat2label = {cat: self.classes.index(cat) for cat in self.classes}
+        self.label2cat = {self.cat2label[t]: t for t in self.cat2label}
+        self.cat_ids = np.array(
+            [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39])
+        self.cat_ids2class = {
+            nyu40id: i
+            for i, nyu40id in enumerate(list(self.cat_ids))
+        }
+        assert split in ['train', 'val', 'test']
+        split_dir = os.path.join(self.root_dir, 'meta_data',
+                                 'scannetv2_%s.txt' % split)
+        self.sample_id_list = [x.strip() for x in open(split_dir).readlines()
+                               ] if os.path.exists(split_dir) else None
+
+    def __len__(self):
+        return len(self.sample_id_list)
+
+    def get_box_label(self, idx):
+        box_file = os.path.join(self.root_dir, 'scannet_train_instance_data',
+                                '%s_bbox.npy' % idx)
+        assert os.path.exists(box_file)
+        return np.load(box_file)
+
+    def get_scannet_infos(self,
+                          num_workers=4,
+                          has_label=True,
+                          sample_id_list=None):
+        import concurrent.futures as futures
+
+        def process_single_scene(sample_idx):
+            print('%s sample_idx: %s' % (self.split, sample_idx))
+            info = dict()
+            pc_info = {'num_features': 6, 'lidar_idx': sample_idx}
+            info['point_cloud'] = pc_info
+
+            if has_label:
+                annotations = {}
+                boxes_with_classes = self.get_box_label(
+                    sample_idx)  # k, 6 + class
+                annotations['gt_num'] = boxes_with_classes.shape[0]
+                if annotations['gt_num'] != 0:
+                    minmax_boxes3d = boxes_with_classes[:, :-1]  # k, 6
+                    classes = boxes_with_classes[:, -1]  # k, 1
+                    annotations['name'] = np.array([
+                        self.label2cat[self.cat_ids2class[classes[i]]]
+                        for i in range(annotations['gt_num'])
+                    ])
+                    annotations['location'] = minmax_boxes3d[:, :3]
+                    annotations['dimensions'] = minmax_boxes3d[:, 3:6]
+                    annotations['gt_boxes_upright_depth'] = minmax_boxes3d
+                    annotations['index'] = np.arange(
+                        annotations['gt_num'], dtype=np.int32)
+                    annotations['class'] = np.array([
+                        self.cat_ids2class[classes[i]]
+                        for i in range(annotations['gt_num'])
+                    ])
+                info['annos'] = annotations
+            return info
+
+        sample_id_list = sample_id_list if sample_id_list is not None \
+            else self.sample_id_list
+        with futures.ThreadPoolExecutor(num_workers) as executor:
+            infos = executor.map(process_single_scene, sample_id_list)
+        return list(infos)

tools/data_converter/sunrgbd_converter.py

+import os
+import pickle
+from pathlib import Path
+
+from tools.data_converter.sunrgbd_data_utils import SUNRGBDData
+
+
+def create_sunrgbd_info_file(data_path,
+                             pkl_prefix='sunrgbd',
+                             save_path=None,
+                             use_v1=False):
+    assert os.path.exists(data_path)
+    if save_path is None:
+        save_path = Path(data_path)
+    else:
+        save_path = Path(save_path)
+    assert os.path.exists(save_path)
+    train_filename = save_path / f'{pkl_prefix}_infos_train.pkl'
+    val_filename = save_path / f'{pkl_prefix}_infos_val.pkl'
+    train_dataset = SUNRGBDData(
+        root_path=data_path, split='train', use_v1=use_v1)
+    val_dataset = SUNRGBDData(root_path=data_path, split='val', use_v1=use_v1)
+    sunrgbd_infos_train = train_dataset.get_sunrgbd_infos(has_label=True)
+    with open(train_filename, 'wb') as f:
+        pickle.dump(sunrgbd_infos_train, f)
+    print('Sunrgbd info train file is saved to %s' % train_filename)
+    sunrgbd_infos_val = val_dataset.get_sunrgbd_infos(has_label=True)
+    with open(val_filename, 'wb') as f:
+        pickle.dump(sunrgbd_infos_val, f)
+    print('Sunrgbd info val file is saved to %s' % val_filename)
+
+
+if __name__ == '__main__':
+    create_sunrgbd_info_file(
+        data_path='./data/sunrgbd/sunrgbd_trainval',
+        save_path='./data/sunrgbd')

tools/data_converter/sunrgbd_data_utils.py

+import os
+
+import cv2
+import numpy as np
+import scipy.io as sio
+
+
+def random_sampling(pc, num_sample, replace=None, return_choices=False):
+    """ Input is NxC, output is num_samplexC
+    """
+    if replace is None:
+        replace = (pc.shape[0] < num_sample)
+    choices = np.random.choice(pc.shape[0], num_sample, replace=replace)
+    if return_choices:
+        return pc[choices], choices
+    else:
+        return pc[choices]
+
+
+class SUNRGBDInstance(object):
+
+    def __init__(self, line):
+        data = line.split(' ')
+        data[1:] = [float(x) for x in data[1:]]
+        self.classname = data[0]
+        self.xmin = data[1]
+        self.ymin = data[2]
+        self.xmax = data[1] + data[3]
+        self.ymax = data[2] + data[4]
+        self.box2d = np.array([self.xmin, self.ymin, self.xmax, self.ymax])
+        self.centroid = np.array([data[5], data[6], data[7]])
+        self.w = data[8]
+        self.l = data[9]  # noqa: E741
+        self.h = data[10]
+        self.orientation = np.zeros((3, ))
+        self.orientation[0] = data[11]
+        self.orientation[1] = data[12]
+        self.heading_angle = -1 * np.arctan2(self.orientation[1],
+                                             self.orientation[0])
+        self.box3d = np.concatenate([
+            self.centroid,
+            np.array([self.l * 2, self.w * 2, self.h * 2, self.heading_angle])
+        ])
+
+
+class SUNRGBDData(object):
+    ''' Load and parse object data '''
+
+    def __init__(self, root_path, split='train', use_v1=False):
+        self.root_dir = root_path
+        self.split = split
+        self.split_dir = os.path.join(root_path)
+        self.classes = [
+            'bed', 'table', 'sofa', 'chair', 'toilet', 'desk', 'dresser',
+            'night_stand', 'bookshelf', 'bathtub'
+        ]
+        self.cat2label = {cat: self.classes.index(cat) for cat in self.classes}
+        self.label2cat = {
+            label: self.classes[label]
+            for label in range(len(self.classes))
+        }
+        assert split in ['train', 'val', 'test']
+        split_dir = os.path.join(self.root_dir, '%s_data_idx.txt' % split)
+        self.sample_id_list = [
+            int(x.strip()) for x in open(split_dir).readlines()
+        ] if os.path.exists(split_dir) else None
+
+        self.image_dir = os.path.join(self.split_dir, 'image')
+        self.calib_dir = os.path.join(self.split_dir, 'calib')
+        self.depth_dir = os.path.join(self.split_dir, 'depth')
+        if use_v1:
+            self.label_dir = os.path.join(self.split_dir, 'label_v1')
+        else:
+            self.label_dir = os.path.join(self.split_dir, 'label')
+
+    def __len__(self):
+        return len(self.sample_id_list)
+
+    def get_image(self, idx):
+        img_filename = os.path.join(self.image_dir, '%06d.jpg' % (idx))
+        return cv2.imread(img_filename)
+
+    def get_image_shape(self, idx):
+        image = self.get_image(idx)
+        return np.array(image.shape[:2], dtype=np.int32)
+
+    def get_depth(self, idx):
+        depth_filename = os.path.join(self.depth_dir, '%06d.mat' % (idx))
+        depth = sio.loadmat(depth_filename)['instance']
+        return depth
+
+    def get_calibration(self, idx):
+        calib_filepath = os.path.join(self.calib_dir, '%06d.txt' % (idx))
+        lines = [line.rstrip() for line in open(calib_filepath)]
+        Rt = np.array([float(x) for x in lines[0].split(' ')])
+        Rt = np.reshape(Rt, (3, 3), order='F')
+        K = np.array([float(x) for x in lines[1].split(' ')])
+        return K, Rt
+
+    def get_label_objects(self, idx):
+        label_filename = os.path.join(self.label_dir, '%06d.txt' % (idx))
+        lines = [line.rstrip() for line in open(label_filename)]
+        objects = [SUNRGBDInstance(line) for line in lines]
+        return objects
+
+    def get_sunrgbd_infos(self,
+                          num_workers=4,
+                          has_label=True,
+                          sample_id_list=None):
+        import concurrent.futures as futures
+
+        def process_single_scene(sample_idx):
+            print('%s sample_idx: %s' % (self.split, sample_idx))
+            # convert depth to points
+            SAMPLE_NUM = 50000
+            pc_upright_depth = self.get_depth(sample_idx)
+            # TODO : sample points in loading process and test
+            pc_upright_depth_subsampled = random_sampling(
+                pc_upright_depth, SAMPLE_NUM)
+            np.savez_compressed(
+                os.path.join(self.root_dir, 'lidar', '%06d.npz' % sample_idx),
+                pc=pc_upright_depth_subsampled)
+
+            info = dict()
+            pc_info = {'num_features': 6, 'lidar_idx': sample_idx}
+            info['point_cloud'] = pc_info
+            img_name = os.path.join(self.image_dir, '%06d.jpg' % (sample_idx))
+            img_path = os.path.join(self.image_dir, img_name)
+            image_info = {
+                'image_idx': sample_idx,
+                'image_shape': self.get_image_shape(sample_idx),
+                'image_path': img_path
+            }
+            info['image'] = image_info
+
+            K, Rt = self.get_calibration(sample_idx)
+            calib_info = {'K': K, 'Rt': Rt}
+            info['calib'] = calib_info
+
+            if has_label:
+                obj_list = self.get_label_objects(sample_idx)
+                annotations = {}
+                annotations['gt_num'] = len([
+                    obj.classname for obj in obj_list
+                    if obj.classname in self.cat2label.keys()
+                ])
+                if annotations['gt_num'] != 0:
+                    annotations['name'] = np.array([
+                        obj.classname for obj in obj_list
+                        if obj.classname in self.cat2label.keys()
+                    ])
+                    annotations['bbox'] = np.concatenate([
+                        obj.box2d.reshape(1, 4) for obj in obj_list
+                        if obj.classname in self.cat2label.keys()
+                    ],
+                                                         axis=0)
+                    annotations['location'] = np.concatenate([
+                        obj.centroid.reshape(1, 3) for obj in obj_list
+                        if obj.classname in self.cat2label.keys()
+                    ],
+                                                             axis=0)
+                    annotations['dimensions'] = 2 * np.array([
+                        [obj.l, obj.h, obj.w] for obj in obj_list
+                        if obj.classname in self.cat2label.keys()
+                    ])  # lhw(depth) format
+                    annotations['rotation_y'] = np.array([
+                        obj.heading_angle for obj in obj_list
+                        if obj.classname in self.cat2label.keys()
+                    ])
+                    annotations['index'] = np.arange(
+                        len(obj_list), dtype=np.int32)
+                    annotations['class'] = np.array([
+                        self.cat2label[obj.classname] for obj in obj_list
+                        if obj.classname in self.cat2label.keys()
+                    ])
+                    annotations['gt_boxes_upright_depth'] = np.stack(
+                        [
+                            obj.box3d for obj in obj_list
+                            if obj.classname in self.cat2label.keys()
+                        ],
+                        axis=0)  # (K,8)
+                info['annos'] = annotations
+            return info
+
+        lidar_save_dir = os.path.join(self.root_dir, 'lidar')
+        if not os.path.exists(lidar_save_dir):
+            os.mkdir(lidar_save_dir)
+        sample_id_list = sample_id_list if \
+            sample_id_list is not None else self.sample_id_list
+        with futures.ThreadPoolExecutor(num_workers) as executor:
+            infos = executor.map(process_single_scene, sample_id_list)
+        return list(infos)