from __future__ import annotations import argparse import json import math import sys from pathlib import Path from typing import Any import numpy as np CODE_ROOT = Path(__file__).resolve().parent SRC_ROOT = CODE_ROOT / "src" if str(SRC_ROOT) not in sys.path: sys.path.insert(0, str(SRC_ROOT)) from physformer.data.obj_io import load_obj_vertices_faces DEFAULT_EVAL_ASSETS_ROOT = CODE_ROOT / "eval_assets" DEFAULT_PRECOMP_ROOT = DEFAULT_EVAL_ASSETS_ROOT / "eval_precomp" DEFAULT_DATA_ROOT = DEFAULT_EVAL_ASSETS_ROOT / "eval_data" DEFAULT_SPLIT_FILE = DEFAULT_EVAL_ASSETS_ROOT / "eval_split.json" DEFAULT_SAMPLE_NAMES = "2obj_elastic,3obj_rigid,4obj_elastic,5obj_rigid" def _safe_float(value: object, default: float) -> float: try: if value is None: return float(default) return float(value) except Exception: return float(default) def _nested_get(data: dict[str, Any], *keys: str) -> Any: cur: Any = data for key in keys: if not isinstance(cur, dict) or key not in cur: return None cur = cur[key] return cur def _object_density(meta: dict[str, Any], obj: dict[str, Any], default_density: float) -> float: for value in ( obj.get("rho"), obj.get("density"), _nested_get(obj, "rigid", "rho"), _nested_get(obj, "material", "rho"), _nested_get(obj, "material", "density"), meta.get("rho"), meta.get("density"), _nested_get(meta, "pbd", "rho"), _nested_get(meta, "pbd", "density"), _nested_get(meta, "rigid", "rho"), _nested_get(meta, "material", "rho"), _nested_get(meta, "material", "density"), ): if value is not None: return _safe_float(value, float(default_density)) return float(default_density) def _mesh_volume(vertices: np.ndarray, faces: np.ndarray) -> float: verts = np.asarray(vertices, dtype=np.float64) faces_i64 = np.asarray(faces, dtype=np.int64) tris = verts[faces_i64] signed_volume = float(np.sum(np.einsum("ij,ij->i", tris[:, 0], np.cross(tris[:, 1], tris[:, 2]))) / 6.0) return abs(signed_volume) def _resolve_mesh_path(mesh_used: str, *, template_dir: Path | None) -> Path: raw = Path(mesh_used) candidates = [raw, CODE_ROOT / raw] if template_dir is not None: candidates.extend([template_dir / raw.name, template_dir / f"{raw.stem}.obj"]) for path in candidates: if path.is_file(): return path raise FileNotFoundError(f"Could not resolve mesh path {mesh_used!r}; tried {', '.join(map(str, candidates))}") def _load_object_masses(meta: dict[str, Any], *, default_density: float, template_dir: Path | None) -> np.ndarray: objects = meta.get("objects", None) if not isinstance(objects, list) or not objects: raise ValueError("metadata.json must contain a non-empty objects list") masses: list[float] = [] for obj_idx, obj_any in enumerate(objects): if not isinstance(obj_any, dict): raise ValueError(f"objects[{obj_idx}] must be a dict") mesh_used = obj_any.get("mesh_used") or obj_any.get("mesh_source") if not isinstance(mesh_used, str) or not mesh_used: raise ValueError(f"objects[{obj_idx}] needs mesh_used or mesh_source") mesh_path = _resolve_mesh_path(mesh_used, template_dir=template_dir) verts, faces = load_obj_vertices_faces(str(mesh_path)) volume = _mesh_volume(verts, faces) scale = _safe_float(obj_any.get("scale"), 1.0) density = _object_density(meta, obj_any, default_density) mass = float(density) * float(volume) * float(scale) ** 3 if not math.isfinite(mass) or mass <= 0.0: raise ValueError(f"Invalid object mass for object {obj_idx}: {mass}") masses.append(mass) return np.asarray(masses, dtype=np.float64) def _object_ids_from_metadata( meta: dict[str, Any], *, total_vertices: int, template_dir: Path | None, ) -> np.ndarray: object_ids = np.full((int(total_vertices),), -1, dtype=np.int64) objects = meta.get("objects", None) if not isinstance(objects, list) or not objects: raise ValueError("metadata.json must contain a non-empty objects list") cursor = 0 for obj_idx, obj_any in enumerate(objects): if not isinstance(obj_any, dict): raise ValueError(f"objects[{obj_idx}] must be a dict") vertex_range = obj_any.get("vertex_range") if isinstance(vertex_range, list) and len(vertex_range) == 2: start, end = int(vertex_range[0]), int(vertex_range[1]) else: count = int(obj_any.get("vertex_count", 0)) if count <= 0: mesh_used = obj_any.get("mesh_used") or obj_any.get("mesh_source") if isinstance(mesh_used, str) and mesh_used: mesh_path = _resolve_mesh_path(mesh_used, template_dir=template_dir) mesh_vertices, _ = load_obj_vertices_faces(str(mesh_path)) count = int(mesh_vertices.shape[0]) if count <= 0: raise ValueError(f"objects[{obj_idx}] needs vertex_range or positive vertex_count") start, end = cursor, cursor + count if start < 0 or end > int(total_vertices) or end <= start: raise ValueError(f"Invalid vertex range for object {obj_idx}: {start}, {end}") object_ids[start:end] = int(obj_idx) cursor = end if np.any(object_ids < 0): missing = int(np.where(object_ids < 0)[0][0]) raise ValueError(f"No object id assigned for vertex {missing}") return object_ids def _load_gt_vertices(meshes_dir: Path) -> np.ndarray: frame_paths = sorted(meshes_dir.glob("combined_frame_*.obj")) if not frame_paths: frame_paths = sorted(meshes_dir.glob("*.obj")) if not frame_paths: raise FileNotFoundError(f"No OBJ frames found in {meshes_dir}") frames: list[np.ndarray] = [] for path in frame_paths: vertices, _ = load_obj_vertices_faces(str(path)) frames.append(vertices.astype(np.float32, copy=False)) first_shape = frames[0].shape for path, vertices in zip(frame_paths, frames): if vertices.shape != first_shape: raise ValueError(f"Frame vertex shape mismatch in {path}: {vertices.shape} != {first_shape}") return np.stack(frames, axis=0).astype(np.float32, copy=False) def _parse_sample_names(value: str) -> list[str]: names = [part.strip() for part in value.split(",") if part.strip()] if not names: raise ValueError("Expected at least one sample name") return names def _relative_to_code(path: Path) -> str: try: return path.resolve().relative_to(CODE_ROOT.resolve()).as_posix() except ValueError: return path.as_posix() def _publish_metadata(sample_name: str, meta: dict[str, Any], object_masses: np.ndarray, object_ids: np.ndarray) -> dict[str, Any]: objects: list[dict[str, Any]] = [] for obj_idx, mass in enumerate(object_masses.tolist()): where = np.where(object_ids == int(obj_idx))[0] vertex_range = [int(where[0]), int(where[-1]) + 1] if where.size else [0, 0] objects.append( { "object_index": int(obj_idx), "mass": float(mass), "vertex_range": vertex_range, } ) return { "sample_name": str(sample_name), "num_objects": int(len(objects)), "frames": int(meta.get("frames", 0) or 0), "dt": float(meta.get("dt", 0.0) or 0.0), "steps_per_frame": int(meta.get("steps_per_frame", 0) or 0), "original_dataset": { "category": meta.get("category"), "sample_index": meta.get("sample_index"), "mode": meta.get("mode"), }, "objects": objects, } def build_argparser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( "Regenerate compact publication loss-evaluation inputs from raw ground-truth sample folders." ) parser.add_argument( "--sample_root", type=Path, default=None, required=True, help=( "Directory containing raw sample folders such as 2obj_elastic and 3obj_rigid. " "Raw samples are not packaged by default; eval_assets/ is already included for normal evaluation." ), ) parser.add_argument("--sample_names", type=str, default=DEFAULT_SAMPLE_NAMES) parser.add_argument("--precomp_root", type=Path, default=DEFAULT_PRECOMP_ROOT) parser.add_argument("--data_root", type=Path, default=DEFAULT_DATA_ROOT) parser.add_argument("--split_file", type=Path, default=DEFAULT_SPLIT_FILE) parser.add_argument("--split_name", type=str, default="test") parser.add_argument("--density", type=float, default=1000.0) parser.add_argument("--mesh_template_dir", type=Path, default=None) parser.add_argument( "--overwrite", action=argparse.BooleanOptionalAction, default=False, help="Regenerate outputs even when eval precomp/data files already exist.", ) return parser def prepare_eval_inputs(args: argparse.Namespace) -> list[str]: args.sample_root = args.sample_root.expanduser().resolve() if not args.sample_root.is_dir(): raise FileNotFoundError(f"--sample_root is not a directory: {args.sample_root}") sample_names = _parse_sample_names(str(args.sample_names)) selectors: list[str] = [] prepared_any = False for sample_name in sample_names: selector = f"{sample_name}:0" selectors.append(selector) out_npz = args.precomp_root / sample_name / "sample_000000.npz" out_sample = args.data_root / sample_name / "sample_000000" out_metadata = out_sample / "metadata.json" if out_npz.is_file() and out_metadata.is_file() and not bool(args.overwrite): print(f"eval assets already present for {selector}; skipping precompute: {out_npz}") continue src_sample_dir = args.sample_root / sample_name metadata_path = src_sample_dir / "metadata.json" meshes_dir = src_sample_dir / "meshes" velocity_path = src_sample_dir / "vertex_velocities" / "combined_frame_000.npy" if not metadata_path.is_file(): raise FileNotFoundError(f"Missing metadata: {metadata_path}") if not meshes_dir.is_dir(): raise FileNotFoundError(f"Missing meshes dir: {meshes_dir}") if not velocity_path.is_file(): raise FileNotFoundError(f"Missing first-frame velocity: {velocity_path}") with metadata_path.open("r", encoding="utf-8") as f: meta = json.load(f) vertices = _load_gt_vertices(meshes_dir) first_vel = np.load(velocity_path).astype(np.float32, copy=False) if first_vel.shape != vertices.shape[1:]: raise ValueError(f"Velocity shape mismatch for {sample_name}: {first_vel.shape} != {vertices.shape[1:]}") object_ids = _object_ids_from_metadata( meta, total_vertices=int(vertices.shape[1]), template_dir=args.mesh_template_dir, ) num_objects = int(meta.get("num_objects", len(meta.get("objects", [])))) object_masses = _load_object_masses( meta, default_density=float(args.density), template_dir=args.mesh_template_dir, ) if int(object_masses.shape[0]) != int(num_objects): raise ValueError(f"num_objects mismatch for {sample_name}: masses={object_masses.shape[0]} meta={num_objects}") out_npz.parent.mkdir(parents=True, exist_ok=True) out_sample.mkdir(parents=True, exist_ok=True) publish_meta = _publish_metadata(sample_name, meta, object_masses, object_ids) with out_metadata.open("w", encoding="utf-8") as f: json.dump(publish_meta, f, indent=2) f.write("\n") np.savez_compressed( out_npz, vertices=vertices, mask=np.ones(vertices.shape[:2], dtype=np.uint8), object_ids=object_ids.astype(np.int64, copy=False), num_objects=np.asarray(num_objects, dtype=np.int64), first_frame_velocity=first_vel, object_masses=object_masses, source_sample_name=np.asarray(sample_name), ) prepared_any = True print(f"prepared {selector}: {out_npz}") payload = { "description": "Official demo four-sample publication evaluation split.", "dataset_root": _relative_to_code(args.data_root), str(args.split_name): selectors, } if not prepared_any and args.split_file.is_file() and not bool(args.overwrite): print(f"split already present; skipping write: {args.split_file}") print(f"precomp_root: {args.precomp_root}") print(f"data_root: {args.data_root}") return selectors args.split_file.parent.mkdir(parents=True, exist_ok=True) with args.split_file.open("w", encoding="utf-8") as f: json.dump(payload, f, indent=2) f.write("\n") print(f"wrote split: {args.split_file}") print(f"precomp_root: {args.precomp_root}") print(f"data_root: {args.data_root}") return selectors def main() -> None: prepare_eval_inputs(build_argparser().parse_args()) if __name__ == "__main__": main()