Spaces:
Running on Zero
Running on Zero
File size: 5,387 Bytes
9550667 | 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 | """Superpoint (over-segmentation) extraction.
Faithful Python/Numba port of ``segment_mesh`` from
https://github.com/kenomo/segmentator (itself the ScanNet ``Segmentator``
Felzenszwalb-Huttenlocher graph segmentation). AQ3D's ScanNet preprocessing
calls it with ``kThresh=0.01, segMinVerts=20``; the network consumes those
superpoints directly, so reproducing this exactly matters far more than it
looks.
The original is a small C++/PyTorch extension that has to be built with CMake
against libtorch, which is not practical inside a Space image, hence the port.
"""
import numpy as np
from numba import njit
@njit(cache=True, nogil=True)
def _segment_graph(ea, eb, w, order, num_vertices, c, seg_min_verts):
parent = np.arange(num_vertices, dtype=np.int64)
rank = np.zeros(num_vertices, dtype=np.int64)
size = np.ones(num_vertices, dtype=np.int64)
threshold = np.full(num_vertices, c, dtype=np.float32)
# --- pass 1: Felzenszwalb-Huttenlocher merging over ascending weights ---
for i in range(order.shape[0]):
e = order[i]
# find(a)
x = ea[e]
y = x
while y != parent[y]:
y = parent[y]
parent[x] = y
a = y
# find(b)
x = eb[e]
y = x
while y != parent[y]:
y = parent[y]
parent[x] = y
b = y
if a != b:
ww = w[e]
if ww <= threshold[a] and ww <= threshold[b]:
# join(a, b)
if rank[a] > rank[b]:
parent[b] = a
size[a] += size[b]
root = a
else:
parent[a] = b
size[b] += size[a]
if rank[a] == rank[b]:
rank[b] += 1
root = b
threshold[root] = ww + c / size[root]
# --- pass 2: absorb segments smaller than seg_min_verts ---
for i in range(order.shape[0]):
e = order[i]
x = ea[e]
y = x
while y != parent[y]:
y = parent[y]
parent[x] = y
a = y
x = eb[e]
y = x
while y != parent[y]:
y = parent[y]
parent[x] = y
b = y
if a != b and (size[a] < seg_min_verts or size[b] < seg_min_verts):
if rank[a] > rank[b]:
parent[b] = a
size[a] += size[b]
else:
parent[a] = b
size[b] += size[a]
if rank[a] == rank[b]:
rank[b] += 1
out = np.empty(num_vertices, dtype=np.int64)
for q in range(num_vertices):
y = q
while y != parent[y]:
y = parent[y]
parent[q] = y
out[q] = y
return out
def unit_face_normals(vertices: np.ndarray, faces: np.ndarray) -> np.ndarray:
v = vertices.astype(np.float64)
n = np.cross(v[faces[:, 1]] - v[faces[:, 0]], v[faces[:, 2]] - v[faces[:, 0]])
length = np.linalg.norm(n, axis=1, keepdims=True)
return np.divide(n, length, out=np.zeros_like(n), where=length > 1e-20)
def segmentator_vertex_normals(vertices: np.ndarray, faces: np.ndarray) -> np.ndarray:
"""Mean of the incident *unit* face normals (the C++ code's running lerp)."""
nf = unit_face_normals(vertices, faces)
idx = faces.reshape(-1)
vals = np.repeat(nf, 3, axis=0)
nv = np.zeros((vertices.shape[0], 3), dtype=np.float64)
for a in range(3):
nv[:, a] = np.bincount(idx, weights=vals[:, a], minlength=vertices.shape[0])
counts = np.bincount(idx, minlength=vertices.shape[0])
nv /= np.maximum(counts, 1)[:, None]
return nv
def segment_mesh(vertices: np.ndarray, faces: np.ndarray, k_thresh: float = 0.01,
seg_min_verts: int = 20) -> np.ndarray:
"""Return a superpoint id per vertex, relabelled to ``0..S-1``."""
vertices = np.ascontiguousarray(vertices, dtype=np.float32)
faces = np.ascontiguousarray(faces, dtype=np.int64)
num_vertices = vertices.shape[0]
normals = segmentator_vertex_normals(vertices, faces)
# edges (i1,i2), (i1,i3), (i3,i2) per face -- exactly as in the C++ source
ea = np.empty(faces.shape[0] * 3, dtype=np.int64)
eb = np.empty(faces.shape[0] * 3, dtype=np.int64)
ea[0::3], eb[0::3] = faces[:, 0], faces[:, 1]
ea[1::3], eb[1::3] = faces[:, 0], faces[:, 2]
ea[2::3], eb[2::3] = faces[:, 2], faces[:, 1]
p = vertices.astype(np.float64)
d = p[eb] - p[ea]
dd = np.linalg.norm(d, axis=1, keepdims=True)
d = np.divide(d, dd, out=np.zeros_like(d), where=dd > 1e-20)
n1, n2 = normals[ea], normals[eb]
dot = (n1 * n2).sum(1)
dot2 = (n2 * d).sum(1)
w = 1.0 - dot
w = np.where(dot2 > 0, w * w, w).astype(np.float32)
order = np.argsort(w, kind="stable").astype(np.int64)
comps = _segment_graph(ea, eb, w, order, num_vertices,
np.float32(k_thresh), np.int64(seg_min_verts))
# remap component ids to 0..S-1 (upstream does the same)
_, remapped = np.unique(comps, return_inverse=True)
return remapped.astype(np.int64).reshape(-1)
def warmup() -> None:
"""Trigger Numba JIT compilation once, at import time."""
v = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0]], dtype=np.float32)
f = np.array([[0, 1, 2], [1, 3, 2]], dtype=np.int64)
segment_mesh(v, f)
|