File size: 1,549 Bytes
424188c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
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')
|