e4cfa9d1e9
Fork of dgrauet/Hunyuan3D-2.1-mlx + our generate_e2e.py CLI, env-tunable remesh (HY3D_REMESH_FACES), and HARDWARE.md. Upstream training data (mini_trainset) and demo images stripped — inference needs none of it. Full upstream history: github.com/dgrauet/Hunyuan3D-2.1-mlx
259 lines
9.1 KiB
Python
259 lines
9.1 KiB
Python
"""Pure-Python port of meshVerticeInpaint (mesh-aware texture inpainting).
|
|
|
|
The PyTorch pipeline uses a C++ extension (mesh_inpaint_processor.cpp) that
|
|
first propagates colors across the mesh's vertex graph, *then* falls back to
|
|
cv2.inpaint. Without step 1, cv2.inpaint must fill huge UV-atlas gaps using
|
|
only 2D neighborhood info and produces noisy bleeding.
|
|
|
|
Algorithm:
|
|
1. Build vertex adjacency graph from face connectivity.
|
|
2. Seed each vertex's color from the texel it lands on (if that texel is
|
|
already painted by back-projection).
|
|
3. For each uncolored vertex, average colors from colored neighbors using
|
|
inverse-square-distance weights. Iterate until convergence.
|
|
4. Rasterize each face into UV space and fill every interior texel with
|
|
barycentric-interpolated vertex color (this is the dense step that the
|
|
vertex-only stamping was missing).
|
|
|
|
The remaining gaps are left to cv2.inpaint.
|
|
"""
|
|
|
|
from typing import Tuple
|
|
|
|
import numpy as np
|
|
|
|
|
|
def _build_vertex_graph(pos_idx: np.ndarray, vtx_num: int) -> list:
|
|
"""For each vertex, list of connected vertex indices (via face edges)."""
|
|
G = [[] for _ in range(vtx_num)]
|
|
for i in range(pos_idx.shape[0]):
|
|
for k in range(3):
|
|
v0 = int(pos_idx[i, k])
|
|
v1 = int(pos_idx[i, (k + 1) % 3])
|
|
G[v0].append(v1)
|
|
return G
|
|
|
|
|
|
def _seed_vertex_colors(
|
|
texture: np.ndarray, mask: np.ndarray, vtx_uv: np.ndarray,
|
|
pos_idx: np.ndarray, uv_idx: np.ndarray, V: int,
|
|
) -> Tuple[np.ndarray, np.ndarray]:
|
|
"""Read each vertex's color from its painted UV texels.
|
|
|
|
A 3D vertex can show up in several UV islands (UV unwrapping splits along
|
|
seams). Each island gives a potentially different color. We average all
|
|
painted readings per 3D vertex instead of taking the last write — last-
|
|
write-wins introduces visible stripes on faces where one UV vertex was
|
|
seeded from view A and another from view B with slightly different
|
|
diffusion output.
|
|
"""
|
|
H, W, C = texture.shape
|
|
|
|
flat_uv_idx = uv_idx.reshape(-1)
|
|
uv_xy = vtx_uv[flat_uv_idx]
|
|
cols = np.clip(np.round(uv_xy[:, 0] * (W - 1)).astype(np.int64), 0, W - 1)
|
|
rows = np.clip(np.round(uv_xy[:, 1] * (H - 1)).astype(np.int64), 0, H - 1)
|
|
|
|
flat_pos = pos_idx.reshape(-1)
|
|
painted = mask[rows, cols] > 0
|
|
colors = texture[rows, cols]
|
|
|
|
valid_v = flat_pos[painted]
|
|
valid_c = colors[painted]
|
|
|
|
vtx_color = np.zeros((V, C), dtype=np.float32)
|
|
vtx_count = np.zeros(V, dtype=np.float32)
|
|
np.add.at(vtx_color, valid_v, valid_c)
|
|
np.add.at(vtx_count, valid_v, 1.0)
|
|
|
|
vtx_mask = (vtx_count > 0).astype(np.float32)
|
|
nonzero = vtx_count > 0
|
|
vtx_color[nonzero] /= vtx_count[nonzero, None]
|
|
return vtx_color, vtx_mask
|
|
|
|
|
|
def _propagate_colors(
|
|
vtx_color: np.ndarray, vtx_mask: np.ndarray,
|
|
vtx_pos: np.ndarray, G: list,
|
|
) -> None:
|
|
"""Smooth uncolored vertices using inverse-square-dist weighted neighbors.
|
|
|
|
Modifies vtx_color, vtx_mask in place.
|
|
"""
|
|
V, C = vtx_color.shape
|
|
smooth_count = 2
|
|
last_uncolored = -1
|
|
uncolored = np.where(vtx_mask == 0)[0].tolist()
|
|
|
|
while smooth_count > 0 and uncolored:
|
|
next_uncolored = []
|
|
for v in uncolored:
|
|
if vtx_mask[v] > 0:
|
|
continue
|
|
p0 = vtx_pos[v]
|
|
acc = np.zeros(C, dtype=np.float32)
|
|
tw = 0.0
|
|
for vn in G[v]:
|
|
if vtx_mask[vn] > 0:
|
|
d = np.linalg.norm(vtx_pos[vn] - p0)
|
|
w = 1.0 / max(d, 1e-4)
|
|
w = w * w
|
|
acc += vtx_color[vn] * w
|
|
tw += w
|
|
if tw > 0:
|
|
vtx_color[v] = acc / tw
|
|
vtx_mask[v] = 1.0
|
|
else:
|
|
next_uncolored.append(v)
|
|
|
|
if len(next_uncolored) == last_uncolored:
|
|
smooth_count -= 1
|
|
else:
|
|
smooth_count += 1
|
|
last_uncolored = len(next_uncolored)
|
|
uncolored = next_uncolored
|
|
|
|
|
|
def _rasterize_face_with_colors(
|
|
tex: np.ndarray, msk: np.ndarray,
|
|
uv0: np.ndarray, uv1: np.ndarray, uv2: np.ndarray,
|
|
c0: np.ndarray, c1: np.ndarray, c2: np.ndarray,
|
|
margin: float = 0.0,
|
|
) -> None:
|
|
"""Fill UNPAINTED texels inside a UV triangle with bary-interpolated colors.
|
|
|
|
Only writes to texels where ``msk == 0``. Texels that were already painted
|
|
by back-projection keep their per-view color — overwriting them with a
|
|
smoothed bary-interp blend reintroduces the noise we're trying to avoid.
|
|
|
|
The optional ``margin`` parameter expands the rasterized region by N
|
|
pixels outside the strict triangle (conservative rasterization). This
|
|
extends each face's color into adjacent gutter texels so mipmap
|
|
downsampling and bilinear sampling at island edges don't pull in colors
|
|
from unrelated UV islands → reduces visible color bleeding on the
|
|
rendered mesh.
|
|
|
|
Modifies tex, msk in place.
|
|
"""
|
|
H, W, _ = tex.shape
|
|
|
|
p0 = np.array([uv0[0] * (W - 1), uv0[1] * (H - 1)], dtype=np.float32)
|
|
p1 = np.array([uv1[0] * (W - 1), uv1[1] * (H - 1)], dtype=np.float32)
|
|
p2 = np.array([uv2[0] * (W - 1), uv2[1] * (H - 1)], dtype=np.float32)
|
|
|
|
pad = int(np.ceil(margin)) + 1
|
|
x_min = max(int(np.floor(min(p0[0], p1[0], p2[0]))) - pad, 0)
|
|
x_max = min(int(np.ceil(max(p0[0], p1[0], p2[0]))) + pad, W - 1)
|
|
y_min = max(int(np.floor(min(p0[1], p1[1], p2[1]))) - pad, 0)
|
|
y_max = min(int(np.ceil(max(p0[1], p1[1], p2[1]))) + pad, H - 1)
|
|
if x_max < x_min or y_max < y_min:
|
|
return
|
|
|
|
xs, ys = np.meshgrid(
|
|
np.arange(x_min, x_max + 1), np.arange(y_min, y_max + 1),
|
|
indexing="xy",
|
|
)
|
|
xs = xs.astype(np.float32) + 0.5
|
|
ys = ys.astype(np.float32) + 0.5
|
|
|
|
denom = (p1[1] - p2[1]) * (p0[0] - p2[0]) + (p2[0] - p1[0]) * (p0[1] - p2[1])
|
|
if abs(denom) < 1e-9:
|
|
return
|
|
inv_denom = 1.0 / denom
|
|
w0 = ((p1[1] - p2[1]) * (xs - p2[0]) + (p2[0] - p1[0]) * (ys - p2[1])) * inv_denom
|
|
w1 = ((p2[1] - p0[1]) * (xs - p2[0]) + (p0[0] - p2[0]) * (ys - p2[1])) * inv_denom
|
|
w2 = 1.0 - w0 - w1
|
|
|
|
# Conservative rasterization: include texels slightly outside the triangle.
|
|
# Negative threshold allows up to `margin / mean_edge_len` slack in bary.
|
|
if margin > 0:
|
|
# Approximate margin in barycentric units: 1 / sqrt(area_in_pixels)
|
|
area = abs(denom) * 0.5
|
|
slack = margin / max(np.sqrt(area), 1.0)
|
|
thresh = -slack
|
|
else:
|
|
thresh = 0.0
|
|
inside = (w0 >= thresh) & (w1 >= thresh) & (w2 >= thresh)
|
|
if not inside.any():
|
|
return
|
|
|
|
msk_window = msk[y_min:y_max + 1, x_min:x_max + 1]
|
|
fill = inside & (msk_window == 0)
|
|
if not fill.any():
|
|
return
|
|
|
|
# Clamp barycentric to [0,1] for color interp outside the strict triangle
|
|
bw0 = np.clip(w0[fill], 0, 1)[:, None]
|
|
bw1 = np.clip(w1[fill], 0, 1)[:, None]
|
|
bw2 = np.clip(w2[fill], 0, 1)[:, None]
|
|
s = bw0 + bw1 + bw2
|
|
bw0 /= s; bw1 /= s; bw2 /= s
|
|
interp = bw0 * c0 + bw1 * c1 + bw2 * c2
|
|
|
|
rows_local, cols_local = np.where(fill)
|
|
rows_global = rows_local + y_min
|
|
cols_global = cols_local + x_min
|
|
tex[rows_global, cols_global] = interp
|
|
msk[rows_global, cols_global] = 255
|
|
|
|
|
|
def mesh_vertex_inpaint(
|
|
texture: np.ndarray,
|
|
mask: np.ndarray,
|
|
vtx_pos: np.ndarray,
|
|
vtx_uv: np.ndarray,
|
|
pos_idx: np.ndarray,
|
|
uv_idx: np.ndarray,
|
|
) -> Tuple[np.ndarray, np.ndarray]:
|
|
"""Propagate colors across the mesh to fill UV gaps.
|
|
|
|
Args:
|
|
texture: (H, W, C) float32 in [0, 1].
|
|
mask: (H, W) uint8, 255 = painted, 0 = unpainted.
|
|
vtx_pos: (V, 3) float32 vertex positions.
|
|
vtx_uv: (V_uv, 2) float32 UV coords (V already flipped by set_mesh).
|
|
pos_idx: (F, 3) int face->vertex indices.
|
|
uv_idx: (F, 3) int face->uv indices.
|
|
|
|
Returns:
|
|
Tuple (new_texture, new_mask).
|
|
"""
|
|
tex = texture.copy()
|
|
msk = mask.copy()
|
|
V = vtx_pos.shape[0]
|
|
|
|
G = _build_vertex_graph(pos_idx, V)
|
|
vtx_color, vtx_mask = _seed_vertex_colors(
|
|
tex, msk, vtx_uv, pos_idx, uv_idx, V,
|
|
)
|
|
_propagate_colors(vtx_color, vtx_mask, vtx_pos, G)
|
|
|
|
# Two passes:
|
|
# 1. Strict bary fill of each face's UV interior (only unpainted texels).
|
|
# 2. Conservative dilation 2px outside each face for mipmap-safe edges.
|
|
for f in range(pos_idx.shape[0]):
|
|
v0, v1, v2 = int(pos_idx[f, 0]), int(pos_idx[f, 1]), int(pos_idx[f, 2])
|
|
if vtx_mask[v0] < 1.0 or vtx_mask[v1] < 1.0 or vtx_mask[v2] < 1.0:
|
|
continue
|
|
u0, u1, u2 = int(uv_idx[f, 0]), int(uv_idx[f, 1]), int(uv_idx[f, 2])
|
|
_rasterize_face_with_colors(
|
|
tex, msk,
|
|
vtx_uv[u0], vtx_uv[u1], vtx_uv[u2],
|
|
vtx_color[v0], vtx_color[v1], vtx_color[v2],
|
|
margin=0.0,
|
|
)
|
|
|
|
for f in range(pos_idx.shape[0]):
|
|
v0, v1, v2 = int(pos_idx[f, 0]), int(pos_idx[f, 1]), int(pos_idx[f, 2])
|
|
if vtx_mask[v0] < 1.0 or vtx_mask[v1] < 1.0 or vtx_mask[v2] < 1.0:
|
|
continue
|
|
u0, u1, u2 = int(uv_idx[f, 0]), int(uv_idx[f, 1]), int(uv_idx[f, 2])
|
|
_rasterize_face_with_colors(
|
|
tex, msk,
|
|
vtx_uv[u0], vtx_uv[u1], vtx_uv[u2],
|
|
vtx_color[v0], vtx_color[v1], vtx_color[v2],
|
|
margin=4.0,
|
|
)
|
|
|
|
return tex, msk
|