|
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') |
|
|