64 lines
2.0 KiB
Python
64 lines
2.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Decimate PLY and SCP to cosma-vm after a job completes."""
|
|
import argparse
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
COSMA_VM = "cosma@192.168.0.83"
|
|
COSMA_DATA = "/data/cosma"
|
|
MAX_PTS = 300_000
|
|
|
|
|
|
def find_ply(frames_dir: Path) -> Path | None:
|
|
for candidate in ["model.ply", "output.ply", "reconstruction.ply"]:
|
|
p = frames_dir / candidate
|
|
if p.exists():
|
|
return p
|
|
plys = list(frames_dir.glob("*.ply"))
|
|
return plys[0] if plys else None
|
|
|
|
|
|
def decimate_ply(src: str, dst: str, max_pts: int = MAX_PTS) -> None:
|
|
import open3d as o3d
|
|
import numpy as np
|
|
src_path = Path(src)
|
|
if not src_path.exists():
|
|
raise FileNotFoundError(src)
|
|
pcd = o3d.io.read_point_cloud(str(src_path))
|
|
n = len(pcd.points)
|
|
if n > max_pts:
|
|
vol = float(np.prod(pcd.get_max_bound() - pcd.get_min_bound()))
|
|
vox = max((vol / max_pts) ** (1 / 3), 0.02)
|
|
pcd = pcd.voxel_down_sample(vox)
|
|
o3d.io.write_point_cloud(dst, pcd)
|
|
print(f"Decimated {n} -> {len(pcd.points)} pts -> {dst}", flush=True)
|
|
|
|
|
|
def main() -> None:
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("job_id", type=int)
|
|
p.add_argument("--frames-dir", required=True)
|
|
p.add_argument("--cosma-vm", default=COSMA_VM)
|
|
p.add_argument("--cosma-data", default=COSMA_DATA)
|
|
args = p.parse_args()
|
|
|
|
frames_dir = Path(args.frames_dir)
|
|
ply_src = find_ply(frames_dir)
|
|
if ply_src is None:
|
|
print(f"No PLY found in {frames_dir}", flush=True)
|
|
sys.exit(0)
|
|
|
|
ply_dec = frames_dir / "model_decimated.ply"
|
|
decimate_ply(str(ply_src), str(ply_dec))
|
|
|
|
remote_dir = args.cosma_data
|
|
subprocess.run(["ssh", args.cosma_vm, f"mkdir -p {remote_dir}"], check=True)
|
|
dst = f"job_{args.job_id}_decimated.ply"
|
|
subprocess.run(["scp", str(ply_dec), f"{args.cosma_vm}:{remote_dir}/{dst}"], check=True)
|
|
print(f"SCP done -> {args.cosma_vm}:{remote_dir}/{dst}", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|