Spaces:
Running on Zero
Running on Zero
| from __future__ import annotations | |
| import os | |
| from itertools import product | |
| from typing import Dict, List, Tuple | |
| import numpy as np | |
| def load_obj_vertices_faces(path: str) -> Tuple[np.ndarray, np.ndarray]: | |
| """ | |
| Minimal OBJ loader (positions + faces). | |
| - Supports `v` and `f` lines. | |
| - Faces are triangulated via fan triangulation. | |
| """ | |
| vertices: List[List[float]] = [] | |
| faces: List[List[int]] = [] | |
| with open(path, "r", encoding="utf-8") as f: | |
| for line in f: | |
| if line.startswith("v "): | |
| parts = line.strip().split() | |
| if len(parts) >= 4: | |
| vertices.append([float(parts[1]), float(parts[2]), float(parts[3])]) | |
| elif line.startswith("f "): | |
| parts = line.strip().split()[1:] | |
| if len(parts) < 3: | |
| continue | |
| face = [] | |
| for p in parts: | |
| idx_str = p.split("/")[0] | |
| if not idx_str: | |
| continue | |
| idx = int(idx_str) | |
| # OBJ indices are 1-based; negatives are relative to end | |
| if idx < 0: | |
| idx = len(vertices) + idx | |
| else: | |
| idx = idx - 1 | |
| face.append(idx) | |
| if len(face) >= 3: | |
| faces.append(face) | |
| if not vertices or not faces: | |
| raise ValueError(f"OBJ has no vertices/faces: {path}") | |
| verts = np.asarray(vertices, dtype=np.float32) | |
| tri_faces = [] | |
| for face in faces: | |
| v0 = face[0] | |
| for i in range(1, len(face) - 1): | |
| tri_faces.append([v0, face[i], face[i + 1]]) | |
| if not tri_faces: | |
| raise ValueError(f"OBJ has no triangulated faces: {path}") | |
| return verts, np.asarray(tri_faces, dtype=np.int64) | |
| def vertices_faces_to_triangles(vertices: np.ndarray, faces: np.ndarray) -> np.ndarray: | |
| tris = vertices[faces] # (T,3,3) | |
| return tris.reshape(tris.shape[0], 9).astype(np.float32) | |
| def load_obj_triangles(path: str) -> np.ndarray: | |
| v, f = load_obj_vertices_faces(path) | |
| return vertices_faces_to_triangles(v, f) | |
| def triangles_to_vertices_faces(triangles: np.ndarray, tolerance: float = 1e-6) -> Tuple[np.ndarray, np.ndarray]: | |
| """ | |
| Converts (T,9) triangles back to a vertex list + face indices by merging near-duplicate vertices. | |
| """ | |
| triangles = np.asarray(triangles, dtype=np.float32) | |
| if triangles.ndim != 2 or triangles.shape[1] != 9: | |
| raise ValueError(f"Expected triangles shape (T,9), got {triangles.shape}") | |
| vertices_list: List[np.ndarray] = [] | |
| faces: List[List[int]] = [] | |
| # Hash buckets of vertices on a tolerance grid to avoid O(N^2) allclose checks. | |
| # We still verify with np.allclose, but only against candidates in neighboring buckets. | |
| if tolerance <= 0: | |
| raise ValueError(f"tolerance must be > 0, got {tolerance}") | |
| inv_tol = 1.0 / float(tolerance) | |
| buckets: Dict[Tuple[int, int, int], List[int]] = {} | |
| neighbor_offsets = list(product((-1, 0, 1), repeat=3)) | |
| def _bucket_key(v: np.ndarray) -> Tuple[int, int, int]: | |
| # floor-based bucket; neighbor search handles boundary cases. | |
| q = np.floor(v * inv_tol).astype(np.int64) | |
| return int(q[0]), int(q[1]), int(q[2]) | |
| def find_vertex_idx(vertex: np.ndarray) -> int: | |
| key = _bucket_key(vertex) | |
| for dx, dy, dz in neighbor_offsets: | |
| cand_key = (key[0] + dx, key[1] + dy, key[2] + dz) | |
| for idx in buckets.get(cand_key, ()): | |
| if np.allclose(vertex, vertices_list[idx], atol=tolerance): | |
| return idx | |
| idx = len(vertices_list) | |
| vertices_list.append(np.array(vertex, dtype=np.float32, copy=True)) | |
| buckets.setdefault(key, []).append(idx) | |
| return idx | |
| for tri in triangles: | |
| tri3 = tri.reshape(3, 3) | |
| face_idx = [find_vertex_idx(v) for v in tri3] | |
| faces.append(face_idx) | |
| vertices = np.stack(vertices_list, axis=0).astype(np.float32) | |
| faces = np.asarray(faces, dtype=np.int64) | |
| return vertices, faces | |
| def save_obj(path: str, triangles: np.ndarray) -> None: | |
| os.makedirs(os.path.dirname(path), exist_ok=True) | |
| v, f = triangles_to_vertices_faces(triangles) | |
| with open(path, "w", encoding="utf-8") as fp: | |
| for vert in v: | |
| fp.write(f"v {vert[0]} {vert[1]} {vert[2]}\n") | |
| for face in f: | |
| # OBJ is 1-indexed | |
| fp.write(f"f {face[0] + 1} {face[1] + 1} {face[2] + 1}\n") | |
| def save_obj_vertices_faces(path: str, vertices: np.ndarray, faces: np.ndarray) -> None: | |
| """ | |
| Saves an OBJ from explicit vertices and faces (triangles). | |
| """ | |
| os.makedirs(os.path.dirname(path), exist_ok=True) | |
| vertices = np.asarray(vertices, dtype=np.float32) | |
| faces = np.asarray(faces, dtype=np.int64) | |
| if vertices.ndim != 2 or vertices.shape[1] != 3: | |
| raise ValueError(f"Expected vertices shape (V,3), got {vertices.shape}") | |
| if faces.ndim != 2 or faces.shape[1] != 3: | |
| raise ValueError(f"Expected faces shape (F,3), got {faces.shape}") | |
| with open(path, "w", encoding="utf-8") as fp: | |
| for vert in vertices: | |
| fp.write(f"v {vert[0]} {vert[1]} {vert[2]}\n") | |
| for face in faces: | |
| fp.write(f"f {int(face[0]) + 1} {int(face[1]) + 1} {int(face[2]) + 1}\n") | |
| def fix_num_triangles( | |
| triangles: np.ndarray, | |
| *, | |
| num_triangles: int, | |
| triangle_sampling: str, | |
| pad_value: float, | |
| sample_idx: np.ndarray | None = None, | |
| ) -> Tuple[np.ndarray, np.ndarray, np.ndarray | None]: | |
| """ | |
| Enforces a fixed number of triangles. | |
| Returns: (triangles_fixed, mask, sample_idx_used) | |
| - mask is (num_triangles,) with 1 for real triangles, 0 for padding. | |
| - if input has more triangles and sample_idx is None, generates sample_idx (for temporal consistency). | |
| """ | |
| t = triangles.shape[0] | |
| if t == num_triangles: | |
| return triangles, np.ones((num_triangles,), dtype=np.float32), sample_idx | |
| if t > num_triangles: | |
| if sample_idx is not None and int(sample_idx.max()) >= t: | |
| # Topology / face count changed across frames; fall back to re-sampling. | |
| sample_idx = None | |
| if sample_idx is None: | |
| if triangle_sampling == "random": | |
| sample_idx = np.random.choice(t, num_triangles, replace=False) | |
| elif triangle_sampling == "first": | |
| sample_idx = np.arange(num_triangles) | |
| else: | |
| raise ValueError(f"Unknown triangle_sampling: {triangle_sampling}") | |
| triangles = triangles[sample_idx] | |
| return triangles, np.ones((num_triangles,), dtype=np.float32), sample_idx | |
| pad = np.full((num_triangles - t, 9), pad_value, dtype=triangles.dtype) | |
| out = np.concatenate([triangles, pad], axis=0) | |
| mask = np.zeros((num_triangles,), dtype=np.float32) | |
| mask[:t] = 1.0 | |
| return out, mask, sample_idx | |
| def fix_num_faces( | |
| faces: np.ndarray, | |
| *, | |
| num_faces: int, | |
| face_sampling: str, | |
| sample_idx: np.ndarray | None = None, | |
| ) -> Tuple[np.ndarray, np.ndarray, np.ndarray | None]: | |
| """ | |
| Enforces a fixed number of faces (triangles) for a connectivity template. | |
| Returns: (faces_fixed, mask, sample_idx_used) | |
| - mask is (num_faces,) with 1 for real faces, 0 for padding. | |
| - sample_idx_used is the indices into the original faces when sub-sampling. | |
| """ | |
| faces = np.asarray(faces, dtype=np.int64) | |
| if faces.ndim != 2 or faces.shape[1] != 3: | |
| raise ValueError(f"Expected faces shape (F,3), got {faces.shape}") | |
| f = faces.shape[0] | |
| if f == num_faces: | |
| return faces, np.ones((num_faces,), dtype=np.float32), sample_idx | |
| if f > num_faces: | |
| if sample_idx is None: | |
| if face_sampling == "random": | |
| sample_idx = np.random.choice(f, num_faces, replace=False) | |
| elif face_sampling == "first": | |
| sample_idx = np.arange(num_faces) | |
| else: | |
| raise ValueError(f"Unknown face_sampling: {face_sampling}") | |
| faces = faces[sample_idx] | |
| return faces, np.ones((num_faces,), dtype=np.float32), sample_idx | |
| pad = np.zeros((num_faces - f, 3), dtype=np.int64) | |
| out = np.concatenate([faces, pad], axis=0) | |
| mask = np.zeros((num_faces,), dtype=np.float32) | |
| mask[:f] = 1.0 | |
| return out, mask, sample_idx | |