File size: 4,029 Bytes
23e73f9 | 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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | import os
import yaml
import subprocess
import time
import shutil
def run_step(name, cmd, env):
print(f" [Step] {name}...")
start_time = time.time()
try:
result = subprocess.run(cmd, shell=True, env=env)
elapsed = time.time() - start_time
if result.returncode != 0:
print(f" [Failed] {name}. Elapsed: {elapsed:.1f}s")
return False
print(f" [Completed] {name}. Elapsed: {elapsed:.1f}s")
return True
except Exception as e:
print(f" [Exception] {name}: {e}")
return False
def main():
config_path = "/root/autodl-tmp/SplatAtlas/configs/2dgs_benchmark.yaml"
with open(config_path, "r") as f:
cfg = yaml.safe_load(f)
method = cfg["method_name"]
iters = cfg["global_settings"]["iterations"]
curr_env = os.environ.copy()
curr_env["PYTHONPATH"] = f"/root/autodl-tmp/SplatAtlas:/root/autodl-tmp/2dgs_official:{curr_env.get('PYTHONPATH', '')}"
for ds in cfg["datasets"]:
print(f"\n>>>> Dataset Domain: {ds['name']} <<<<")
for scene in ds["scenes"]:
source = os.path.join(ds["base_path"], scene)
model_out = f"/root/autodl-tmp/SplatAtlas/outputs/{method}_{scene}"
res = ds.get("resolution", 1)
if not os.path.exists(source):
continue
print(f"\n[{method.upper()}] Processing Scene: {scene} (r={res})")
ply_path = os.path.join(model_out, f"point_cloud/iteration_{iters}/point_cloud.ply")
metrics_json = os.path.join(model_out, f"metrics_test_iter{iters}.json")
render_flag = os.path.join(model_out, f"render_complete_{iters}.flag")
renders_dir = os.path.join(model_out, f"renders_test_{iters}")
depths_dir = os.path.join(model_out, f"depths_test_{iters}")
if os.path.exists(ply_path):
print(f" [Checkpoint] Ply file found.")
if not os.path.exists(render_flag):
print(f" [Cleanup] Previous render incomplete, re-rendering...")
for d in ["renders_test", "renders_train", "depths_test", "depths_train", "normals_test", "gt_test"]:
d_path = os.path.join(model_out, f"{d}_{iters}")
if os.path.exists(d_path): shutil.rmtree(d_path)
render_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_render.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters} --resolution {res}"
if not run_step("Offline Rendering", render_cmd, curr_env): continue
else:
print(f" [Checkpoint] Render assets complete.")
else:
train_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/main_train.py --method {method} --source_path {source} --model_path {model_out} --iterations {iters} --resolution {res} --track_decoupling"
if not run_step("Training", train_cmd, curr_env): continue
if not os.path.exists(metrics_json):
eval_cmd = f"python /root/autodl-tmp/SplatAtlas/ufd_evalkit/run_eval.py --method {method} --scene {scene} --render_dir {renders_dir} --gt_dir {model_out}/gt_test_{iters} --ply_path {ply_path} --output_json {metrics_json} --colmap_dir {source} --depth_dir {depths_dir}"
run_step("Metrics Calculation", eval_cmd, curr_env)
else:
print(f" [Checkpoint] Metrics JSON already exists.")
physics_cmd = f"python /root/autodl-tmp/SplatAtlas/scripts/compute_offline_physics.py --method {method} --source_path {source} --model_path {model_out} --iteration {iters}"
run_step("Offline Physics", physics_cmd, curr_env)
run_step("Grand Table", "python /root/autodl-tmp/SplatAtlas/scripts/build_grand_table.py", curr_env)
print(f"Done: {scene}\n")
if __name__ == "__main__":
main()
|