Spaces:
Running on Zero
Running on Zero
File size: 6,825 Bytes
3c58630 | 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 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 | from __future__ import annotations
import json
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Optional
@dataclass(frozen=True)
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")
|