"vscode:/vscode.git/clone" did not exist on "38255d2af15932150606e19bea8200b386cfd36d"
batch_load_scannet_data.py 4.28 KB
Newer Older
1
# Modified from Facebook, Inc. and its affiliates.
liyinhao's avatar
liyinhao committed
2
3
4
5
6
7
8
9
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
""" Batch mode in loading Scannet scenes with vertices and ground truth labels
for semantic and instance segmentations

Usage example: python ./batch_load_scannet_data.py
"""
10
import argparse
liyinhao's avatar
liyinhao committed
11
12
import datetime
import os
13
import os.path as osp
liyinhao's avatar
liyinhao committed
14
15
16
17
18
19
20
21
22
23

import numpy as np
from load_scannet_data import export

SCANNET_DIR = 'scans'
DONOTCARE_CLASS_IDS = np.array([])
OBJ_CLASS_IDS = np.array(
    [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39])


24
25
def export_one_scan(scan_name, output_filename_prefix, max_num_point,
                    label_map_file, scannet_dir):
26
27
28
29
30
31
32
    mesh_file = osp.join(scannet_dir, scan_name, scan_name + '_vh_clean_2.ply')
    agg_file = osp.join(scannet_dir, scan_name,
                        scan_name + '.aggregation.json')
    seg_file = osp.join(scannet_dir, scan_name,
                        scan_name + '_vh_clean_2.0.010000.segs.json')
    # includes axisAlignment info for the train set scans.
    meta_file = osp.join(scannet_dir, scan_name, f'{scan_name}.txt')
liyinhao's avatar
liyinhao committed
33
34
    mesh_vertices, semantic_labels, instance_labels, instance_bboxes, \
        instance2semantic = export(mesh_file, agg_file, seg_file,
35
                                   meta_file, label_map_file, None)
liyinhao's avatar
liyinhao committed
36
37
38
39
40
41
42

    mask = np.logical_not(np.in1d(semantic_labels, DONOTCARE_CLASS_IDS))
    mesh_vertices = mesh_vertices[mask, :]
    semantic_labels = semantic_labels[mask]
    instance_labels = instance_labels[mask]

    num_instances = len(np.unique(instance_labels))
43
    print(f'Num of instances: {num_instances}')
liyinhao's avatar
liyinhao committed
44
45
46
47
48
49

    bbox_mask = np.in1d(instance_bboxes[:, -1], OBJ_CLASS_IDS)
    instance_bboxes = instance_bboxes[bbox_mask, :]
    print('Num of care instances: ', instance_bboxes.shape[0])

    N = mesh_vertices.shape[0]
50
51
    if N > max_num_point:
        choices = np.random.choice(N, max_num_point, replace=False)
liyinhao's avatar
liyinhao committed
52
53
54
55
        mesh_vertices = mesh_vertices[choices, :]
        semantic_labels = semantic_labels[choices]
        instance_labels = instance_labels[choices]

56
57
58
59
    np.save(f'{output_filename_prefix}_vert.npy', mesh_vertices)
    np.save(f'{output_filename_prefix}_sem_label.npy', semantic_labels)
    np.save(f'{output_filename_prefix}_ins_label.npy', instance_labels)
    np.save(f'{output_filename_prefix}_bbox.npy', instance_bboxes)
liyinhao's avatar
liyinhao committed
60
61


62
63
64
def batch_export(max_num_point, output_folder, train_scan_names_file,
                 label_map_file, scannet_dir):
    if not os.path.exists(output_folder):
65
        print(f'Creating new data folder: {output_folder}')
66
        os.mkdir(output_folder)
liyinhao's avatar
liyinhao committed
67

68
69
    train_scan_names = [line.rstrip() for line in open(train_scan_names_file)]
    for scan_name in train_scan_names:
liyinhao's avatar
liyinhao committed
70
71
72
        print('-' * 20 + 'begin')
        print(datetime.datetime.now())
        print(scan_name)
73
74
        output_filename_prefix = osp.join(output_folder, scan_name)
        if osp.isfile(f'{output_filename_prefix}_vert.npy'):
liyinhao's avatar
liyinhao committed
75
76
77
78
            print('File already exists. skipping.')
            print('-' * 20 + 'done')
            continue
        try:
79
80
            export_one_scan(scan_name, output_filename_prefix, max_num_point,
                            label_map_file, scannet_dir)
liyinhao's avatar
liyinhao committed
81
82
83
84
85
        except Exception:
            print('Failed export scan: %s' % (scan_name))
        print('-' * 20 + 'done')


86
87
88
89
90
91
92
93
94
95
96
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--max_num_point',
        default=50000,
        help='The maximum number of the points.')
    parser.add_argument(
        '--output_folder',
        default='./scannet_train_instance_data',
        help='output folder of the result.')
    parser.add_argument(
yinchimaoliang's avatar
yinchimaoliang committed
97
        '--scannet_dir', default='scans', help='scannet data directory.')
98
99
100
101
102
103
104
105
106
107
108
109
110
111
    parser.add_argument(
        '--label_map_file',
        default='meta_data/scannetv2-labels.combined.tsv',
        help='The path of label map file.')
    parser.add_argument(
        '--train_scan_names_file',
        default='meta_data/scannet_train.txt',
        help='The path of the file that stores the scan names.')
    args = parser.parse_args()
    batch_export(args.max_num_point, args.output_folder,
                 args.train_scan_names_file, args.label_map_file,
                 args.scannet_dir)


liyinhao's avatar
liyinhao committed
112
if __name__ == '__main__':
113
    main()