nbv_reconstruction/preprocess/pack_preprocessed_data.py
2024-10-26 03:24:01 -05:00

49 lines
1.9 KiB
Python

import os
import shutil
def pack_scene_data(root, scene, output_dir):
scene_dir = os.path.join(output_dir, scene)
if not os.path.exists(scene_dir):
os.makedirs(scene_dir)
pts_dir = os.path.join(root, scene, "pts")
if os.path.exists(pts_dir):
shutil.move(pts_dir, os.path.join(scene_dir, "pts"))
scan_points_indices_dir = os.path.join(root, scene, "scan_points_indices")
if os.path.exists(scan_points_indices_dir):
shutil.move(scan_points_indices_dir, os.path.join(scene_dir, "scan_points_indices"))
scan_points_file = os.path.join(root, scene, "scan_points.txt")
if os.path.exists(scan_points_file):
shutil.move(scan_points_file, os.path.join(scene_dir, "scan_points.txt"))
model_pts_nrm_file = os.path.join(root, scene, "points_and_normals.txt")
if os.path.exists(model_pts_nrm_file):
shutil.move(model_pts_nrm_file, os.path.join(scene_dir, "points_and_normals.txt"))
camera_dir = os.path.join(root, scene, "camera_params")
if os.path.exists(camera_dir):
shutil.move(camera_dir, os.path.join(scene_dir, "camera_params"))
scene_info_file = os.path.join(root, scene, "scene_info.json")
if os.path.exists(scene_info_file):
shutil.move(scene_info_file, os.path.join(scene_dir, "scene_info.json"))
def pack_all_scenes(root, scene_list, output_dir):
for idx, scene in enumerate(scene_list):
print(f"正在打包场景 {scene} ({idx+1}/{len(scene_list)})")
pack_scene_data(root, scene, output_dir)
if __name__ == "__main__":
root = r"H:\AI\Datasets\nbv_rec_part2"
output_dir = r"H:\AI\Datasets\scene_info_part2"
scene_list = os.listdir(root)
from_idx = 0
to_idx = len(scene_list)
print(f"正在打包场景 {scene_list[from_idx:to_idx]}")
pack_all_scenes(root, scene_list[from_idx:to_idx], output_dir)
print("打包完成")