| import os | |
| from shutil import copyfile | |
| data_base = './montefloor_data/' | |
| dir_names = list(sorted(os.listdir(data_base))) | |
| out_path = './s3d_floorplan' | |
| wrong_s3d_annotations_list = [3261, 3271, 3276, 3296, 3342, 3387, 3398, 3466, 3496] | |
| train_list = [] | |
| val_list = [] | |
| test_list = [] | |
| for dir_name in dir_names: | |
| data_dir = os.path.join(data_base, dir_name) | |
| annot_path = os.path.join(data_dir, 'annot.npy') | |
| if not os.path.exists(annot_path): | |
| continue | |
| data_id = int(dir_name[-5:]) | |
| if data_id in wrong_s3d_annotations_list: | |
| continue | |
| annot_dst = os.path.join(out_path, 'annot', dir_name[-5:] + '.npy') | |
| density_dst = os.path.join(out_path, 'density', dir_name[-5:] + '.png') | |
| normal_dst = os.path.join(out_path, 'normals', dir_name[-5:] + '.png') | |
| density_src = os.path.join(data_dir, 'density.png') | |
| normal_src = os.path.join(data_dir, 'normals.png') | |
| copyfile(normal_src, normal_dst) | |
| copyfile(density_src, density_dst) | |
| copyfile(annot_path, annot_dst) | |
| if 0 <= data_id < 3000: | |
| train_list.append(dir_name[-5:]) | |
| elif data_id < 3250: | |
| val_list.append(dir_name[-5:]) | |
| else: | |
| test_list.append(dir_name[-5:]) | |
| with open(os.path.join(out_path, 'train_list.txt'), 'w') as f: | |
| for item in train_list: | |
| f.write(item + '\n') | |
| with open(os.path.join(out_path, 'valid_list.txt'), 'w') as f: | |
| for item in val_list: | |
| f.write(item + '\n') | |
| with open(os.path.join(out_path, 'test_list.txt'), 'w') as f: | |
| for item in test_list: | |
| f.write(item + '\n') | |