Commit ab96e4f0 authored by liyinhao's avatar liyinhao
Browse files

add empty lines, delete main

parent 6f22a6d3
......@@ -23,8 +23,8 @@ def create_indoor_info_file(data_path,
assert pkl_prefix in ['sunrgbd', 'scannet']
save_path = data_path if save_path is None else save_path
assert os.path.exists(save_path)
train_filename = os.path.join(save_path, f'{pkl_prefix}_infos_train.pkl')
train_filename = os.path.join(save_path, f'{pkl_prefix}_infos_train.pkl')
val_filename = os.path.join(save_path, f'{pkl_prefix}_infos_val.pkl')
if pkl_prefix == 'sunrgbd':
train_dataset = SUNRGBDData(
......@@ -34,22 +34,11 @@ def create_indoor_info_file(data_path,
else:
train_dataset = ScanNetData(root_path=data_path, split='train')
val_dataset = ScanNetData(root_path=data_path, split='val')
infos_train = train_dataset.get_infos(has_label=True)
infos_train = train_dataset.get_infos(has_label=True)
mmcv.dump(infos_train, train_filename, 'pkl')
print(f'{pkl_prefix} info train file is saved to {train_filename}')
infos_val = val_dataset.get_infos(has_label=True)
infos_val = val_dataset.get_infos(has_label=True)
mmcv.dump(infos_val, val_filename, 'pkl')
print(f'{pkl_prefix} info val file is saved to {val_filename}')
if __name__ == '__main__':
create_indoor_info_file(
data_path='./data/scannet',
pkl_prefix='scannet',
save_path='./data/scannet')
create_indoor_info_file(
data_path='./data/sunrgbd/sunrgbd_trainval',
pkl_prefix='sunrgbd',
save_path='./data/sunrgbd')
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