Spaces:
Sleeping
Sleeping
File size: 8,439 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 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 | 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
|