Spaces:
Running on Zero
Running on Zero
File size: 1,763 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 | from __future__ import annotations
from typing import Optional, Tuple
import numpy as np
def fix_num_vertices(
vertices: np.ndarray,
*,
num_vertices: int,
vertex_sampling: str,
pad_value: float,
sample_idx: Optional[np.ndarray] = None,
) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
"""Pad or sample a vertex array to a fixed token count.
Returns:
vertices_fixed: Array with shape ``(num_vertices, 3)``.
mask: Float mask with shape ``(num_vertices,)``; real vertices are 1 and padding is 0.
sample_idx: Reusable sampled indices when the input was downsampled.
"""
target_vertices = int(num_vertices)
current_vertices = int(vertices.shape[0])
if current_vertices == target_vertices:
return vertices, np.ones((target_vertices,), dtype=np.float32), sample_idx
if current_vertices > target_vertices:
if sample_idx is not None and int(sample_idx.max()) >= current_vertices:
sample_idx = None
if sample_idx is None:
if vertex_sampling == "random":
sample_idx = np.random.choice(current_vertices, target_vertices, replace=False)
elif vertex_sampling == "first":
sample_idx = np.arange(target_vertices)
else:
raise ValueError(f"Unknown vertex_sampling: {vertex_sampling}")
return vertices[sample_idx], np.ones((target_vertices,), dtype=np.float32), sample_idx
pad = np.full((target_vertices - current_vertices, 3), pad_value, dtype=vertices.dtype)
vertices_fixed = np.concatenate([vertices, pad], axis=0)
mask = np.zeros((target_vertices,), dtype=np.float32)
mask[:current_vertices] = 1.0
return vertices_fixed, mask, sample_idx
|