Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| import json | |
| import os | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| from typing import Optional | |
| class MultiObjSceneInfo: | |
| mesh_paths: list[str] | |
| mesh_names: list[str] | |
| vertex_counts: list[int] | |
| vertex_slices: list[tuple[int, int]] # (start,end) into the combined vertex array | |
| total_vertices: int | |
| def default_vertex_count_json_path() -> str: | |
| src_root = Path(__file__).resolve().parents[2] | |
| code_root = src_root.parent | |
| candidates = [ | |
| src_root / "official_demo_inference" / "configs" / "vertex_counts_multiobj_all.json", | |
| code_root / "configs" / "vertex_counts_multiobj_all.json", | |
| src_root / "mesh_primitives" / "vertex_counts_multiobj_all.json", | |
| ] | |
| for path in candidates: | |
| if path.is_file(): | |
| return str(path) | |
| return str(candidates[0]) | |
| def load_mesh_vertex_counts(json_path: str) -> dict[str, int]: | |
| with open(json_path, "r", encoding="utf-8") as f: | |
| data = json.load(f) | |
| if isinstance(data, dict) and "counts" in data and isinstance(data["counts"], dict): | |
| counts = data["counts"] | |
| elif isinstance(data, dict): | |
| # Back-compat: allow direct mapping dict[str,int]. | |
| counts = data | |
| else: | |
| raise ValueError(f"Unsupported JSON format in {json_path}: expected dict, got {type(data)}") | |
| out: dict[str, int] = {} | |
| for k, v in counts.items(): | |
| if not isinstance(k, str): | |
| k = str(k) | |
| if not isinstance(v, (int, float)) or int(v) <= 0: | |
| raise ValueError(f"Invalid vertex count for key='{k}' in {json_path}: {v}") | |
| out[k] = int(v) | |
| return out | |
| def _mesh_key_candidates(mesh_path: str) -> list[str]: | |
| p = Path(str(mesh_path)) | |
| stem = p.stem | |
| name = p.name | |
| # Some metadata stores absolute paths; some users prefer stem keys. | |
| # Try several candidates (and their lowercase forms). | |
| cands = [stem, name, str(p)] | |
| out: list[str] = [] | |
| for c in cands: | |
| if c and c not in out: | |
| out.append(c) | |
| lc = c.lower() | |
| if lc and lc not in out: | |
| out.append(lc) | |
| return out | |
| def vertex_count_for_mesh(mesh_path: str, vertex_counts: dict[str, int]) -> int: | |
| for k in _mesh_key_candidates(mesh_path): | |
| if k in vertex_counts: | |
| return int(vertex_counts[k]) | |
| raise KeyError( | |
| f"Mesh '{mesh_path}' not found in vertex count map. Tried keys: {', '.join(_mesh_key_candidates(mesh_path))}" | |
| ) | |
| def _load_metadata(meta_path: str) -> dict: | |
| with open(meta_path, "r", encoding="utf-8") as f: | |
| meta = json.load(f) | |
| if not isinstance(meta, dict): | |
| raise ValueError(f"metadata.json must contain a dict, got {type(meta)}: {meta_path}") | |
| return meta | |
| def scene_info_from_metadata_dict( | |
| meta: dict, | |
| *, | |
| vertex_counts: dict[str, int], | |
| max_num_objects: int, | |
| source: str = "<metadata>", | |
| ) -> MultiObjSceneInfo: | |
| objects = meta.get("objects", None) | |
| if not isinstance(objects, list) or not objects: | |
| raise ValueError(f"metadata missing non-empty 'objects' list: {source}") | |
| mesh_paths: list[str] = [] | |
| mesh_names: list[str] = [] | |
| counts: list[int] = [] | |
| for obj in objects: | |
| if not isinstance(obj, dict): | |
| raise ValueError(f"metadata 'objects' entries must be dicts, got {type(obj)}: {source}") | |
| mesh_used = obj.get("mesh_used") or obj.get("mesh_source") | |
| if not isinstance(mesh_used, str) or not mesh_used: | |
| raise ValueError(f"metadata object missing mesh_used/mesh_source string: {source}") | |
| mesh_paths.append(mesh_used) | |
| mesh_names.append(Path(mesh_used).stem) | |
| # Prefer per-object vertex_count stored in metadata when available. | |
| # This makes inference/precompute robust to new primitive sets without having to regenerate | |
| # a vertex_counts JSON, while still validating against the JSON when possible. | |
| meta_vc = obj.get("vertex_count", None) | |
| meta_vc_int = int(meta_vc) if isinstance(meta_vc, (int, float)) and int(meta_vc) > 0 else None | |
| try: | |
| vc = vertex_count_for_mesh(mesh_used, vertex_counts) | |
| except KeyError: | |
| if meta_vc_int is None: | |
| raise | |
| vc = int(meta_vc_int) | |
| else: | |
| if meta_vc_int is not None and int(meta_vc_int) != int(vc): | |
| raise ValueError( | |
| f"metadata vertex_count={int(meta_vc_int)} disagrees with vertex_counts map={int(vc)} for mesh_used={mesh_used!r}. " | |
| f"meta={source}" | |
| ) | |
| counts.append(int(vc)) | |
| num_obj = len(counts) | |
| if int(num_obj) > int(max_num_objects): | |
| raise ValueError( | |
| f"Scene has num_objects={num_obj} but max_num_objects={int(max_num_objects)}. " | |
| f"Increase --max_num_objects. meta={source}" | |
| ) | |
| slices: list[tuple[int, int]] = [] | |
| cur = 0 | |
| for v in counts: | |
| start = cur | |
| cur += int(v) | |
| slices.append((start, cur)) | |
| return MultiObjSceneInfo( | |
| mesh_paths=mesh_paths, | |
| mesh_names=mesh_names, | |
| vertex_counts=counts, | |
| vertex_slices=slices, | |
| total_vertices=int(cur), | |
| ) | |
| def scene_info_from_metadata( | |
| meta_path: str, | |
| *, | |
| vertex_counts: dict[str, int], | |
| max_num_objects: int, | |
| ) -> MultiObjSceneInfo: | |
| meta = _load_metadata(meta_path) | |
| return scene_info_from_metadata_dict( | |
| meta, | |
| vertex_counts=vertex_counts, | |
| max_num_objects=max_num_objects, | |
| source=meta_path, | |
| ) | |
| def object_ids_from_scene_info(scene: MultiObjSceneInfo) -> list[int]: | |
| ids: list[int] = [] | |
| for obj_id, v in enumerate(scene.vertex_counts): | |
| ids.extend([int(obj_id)] * int(v)) | |
| if len(ids) != int(scene.total_vertices): | |
| raise RuntimeError("object_ids_from_scene_info internal size mismatch") | |
| return ids | |
| def resolve_metadata_path_from_sequence_dir(sequence_dir: str, *, metadata_filename: str = "metadata.json") -> str: | |
| """ | |
| Given a frame directory like <sample_dir>/meshes, return <sample_dir>/metadata.json. | |
| """ | |
| sample_dir = os.path.dirname(str(sequence_dir).rstrip(os.sep)) | |
| return os.path.join(sample_dir, str(metadata_filename)) | |
| def resolve_sample_dir_from_first_frame_obj(obj_path: str) -> str: | |
| """ | |
| Expected layout: | |
| <sample_dir>/meshes/combined_frame_000.obj | |
| """ | |
| frame_dir = os.path.dirname(str(obj_path)) | |
| return os.path.dirname(frame_dir) | |
| def resolve_velocity_path_from_first_frame_obj(obj_path: str, *, velocity_dirname: str) -> str: | |
| sample_dir = resolve_sample_dir_from_first_frame_obj(obj_path) | |
| stem = os.path.splitext(os.path.basename(obj_path))[0] | |
| return os.path.join(sample_dir, str(velocity_dirname), f"{stem}.npy") | |