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
1138 lines
44 KiB
Python
1138 lines
44 KiB
Python
"""MLX mesh renderer for Hunyuan3D texture pipeline.
|
||
|
||
Port of MeshRender's rendering and texture baking pipeline to MLX,
|
||
using mlx_arsenal.rasterize as the Metal rasterization backend.
|
||
"""
|
||
|
||
import mlx.core as mx
|
||
import numpy as np
|
||
|
||
from mlx_arsenal.rasterize import interpolate as mlx_interpolate
|
||
from mlx_arsenal.rasterize import rasterize_triangles
|
||
|
||
from .camera_utils_mlx import (
|
||
get_mv_matrix,
|
||
get_orthographic_projection_matrix,
|
||
get_perspective_projection_matrix,
|
||
transform_pos,
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Rasterizer adapter
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class MLXRasterizer:
|
||
"""Adapter wrapping mlx_arsenal.rasterize to match MeshRender's interface."""
|
||
|
||
@staticmethod
|
||
def rasterize(pos, tri, resolution, clamp_depth=None, use_depth_prior=0):
|
||
"""Rasterize clip-space triangles.
|
||
|
||
Args:
|
||
pos: (1, N, 4) or (N, 4) mx.array clip-space vertices.
|
||
tri: (F, 3) mx.array int32 face indices.
|
||
resolution: (height, width) tuple.
|
||
|
||
Returns:
|
||
findices: (H, W) int32, 1-indexed face IDs.
|
||
barycentric: (H, W, 3) float32.
|
||
"""
|
||
vertices = pos[0] if pos.ndim == 3 else pos
|
||
height, width = int(resolution[0]), int(resolution[1])
|
||
|
||
depth_prior = None
|
||
if clamp_depth is not None and use_depth_prior:
|
||
depth_prior = clamp_depth
|
||
|
||
findices, barycentric = rasterize_triangles(
|
||
vertices.astype(mx.float32),
|
||
tri.astype(mx.int32),
|
||
width,
|
||
height,
|
||
depth_prior=depth_prior,
|
||
)
|
||
return findices, barycentric
|
||
|
||
@staticmethod
|
||
def interpolate(col, findices, barycentric, tri):
|
||
"""Interpolate per-vertex attributes at rasterized pixels.
|
||
|
||
Args:
|
||
col: (1, N, C) or (N, C) mx.array per-vertex data.
|
||
findices: (H, W) int32.
|
||
barycentric: (H, W, 3) float32.
|
||
tri: (F, 3) int32.
|
||
|
||
Returns:
|
||
(1, H, W, C) mx.array.
|
||
"""
|
||
attributes = col[0] if col.ndim == 3 else col
|
||
result = mlx_interpolate(
|
||
attributes.astype(mx.float32),
|
||
findices,
|
||
barycentric,
|
||
tri.astype(mx.int32),
|
||
)
|
||
return mx.expand_dims(result, axis=0)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Mesh renderer
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
class MeshRenderMLX:
|
||
"""MLX mesh renderer for normal / position / alpha rendering.
|
||
|
||
Drop-in replacement for MeshRender's core render pipeline.
|
||
All internal computation uses MLX arrays; public render methods
|
||
return numpy arrays for downstream compatibility (PIL, cv2, etc.).
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
camera_distance=1.45,
|
||
camera_type="orth",
|
||
default_resolution=1024,
|
||
texture_size=1024,
|
||
shader_type="face",
|
||
scale_factor=1.15,
|
||
bake_angle_thres=75,
|
||
boundary_scale=2,
|
||
bake_mode="back_sample",
|
||
raster_mode="mlx",
|
||
):
|
||
self.camera_distance = camera_distance
|
||
self.shader_type = shader_type
|
||
self._default_scale_factor = scale_factor
|
||
self.bake_angle_thres = bake_angle_thres
|
||
self.raster = MLXRasterizer()
|
||
|
||
if isinstance(default_resolution, int):
|
||
self.default_resolution = (default_resolution, default_resolution)
|
||
else:
|
||
self.default_resolution = tuple(default_resolution)
|
||
|
||
if isinstance(texture_size, int):
|
||
self.texture_size = (texture_size, texture_size)
|
||
else:
|
||
self.texture_size = tuple(texture_size)
|
||
|
||
self.bake_unreliable_kernel_size = int(
|
||
(boundary_scale / 512)
|
||
* max(self.default_resolution[0], self.default_resolution[1])
|
||
)
|
||
|
||
if camera_type == "orth":
|
||
self.set_orth_scale(1.2)
|
||
elif camera_type == "perspective":
|
||
self.camera_proj_mat = get_perspective_projection_matrix(
|
||
49.13,
|
||
self.default_resolution[1] / self.default_resolution[0],
|
||
0.01,
|
||
100.0,
|
||
)
|
||
else:
|
||
raise ValueError(f"Unknown camera type: {camera_type}")
|
||
|
||
# Geometry (set by set_mesh)
|
||
self.vtx_pos = None
|
||
self.pos_idx = None
|
||
self.vtx_uv = None
|
||
self.uv_idx = None
|
||
self.scale_factor = 1.0
|
||
self.mesh_normalize_scale_factor = 1.0
|
||
self.mesh_normalize_scale_center = np.array([[0, 0, 0]])
|
||
|
||
# Texture-space data (set by extract_textiles)
|
||
self.tex_position = None
|
||
self.tex_normal = None
|
||
self.tex_grid = None
|
||
self.texture_indices = None
|
||
|
||
# ------------------------------------------------------------------
|
||
# Projection
|
||
# ------------------------------------------------------------------
|
||
|
||
def set_orth_scale(self, ortho_scale):
|
||
self.ortho_scale = ortho_scale
|
||
self.camera_proj_mat = get_orthographic_projection_matrix(
|
||
left=-ortho_scale * 0.5,
|
||
right=ortho_scale * 0.5,
|
||
bottom=-ortho_scale * 0.5,
|
||
top=ortho_scale * 0.5,
|
||
near=0.1,
|
||
far=100,
|
||
)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Mesh loading
|
||
# ------------------------------------------------------------------
|
||
|
||
def set_mesh(self, vtx_pos, pos_idx, vtx_uv=None, uv_idx=None,
|
||
scale_factor=None, auto_center=True):
|
||
"""Load mesh geometry from numpy arrays.
|
||
|
||
Applies the same coordinate transform and normalization
|
||
as the original MeshRender.set_mesh().
|
||
"""
|
||
if scale_factor is None:
|
||
scale_factor = self._default_scale_factor
|
||
|
||
vp = np.array(vtx_pos, dtype=np.float32)
|
||
pi = np.array(pos_idx, dtype=np.int32)
|
||
|
||
# Coordinate transform: flip X/Y, swap Y/Z
|
||
vp[:, [0, 1]] = -vp[:, [0, 1]]
|
||
vp[:, [1, 2]] = vp[:, [2, 1]].copy()
|
||
|
||
if auto_center:
|
||
max_bb = vp.max(axis=0)
|
||
min_bb = vp.min(axis=0)
|
||
center = (max_bb + min_bb) / 2.0
|
||
scale = np.linalg.norm(vp - center, axis=1).max() * 2.0
|
||
vp = (vp - center) * (scale_factor / scale)
|
||
self.scale_factor = scale_factor
|
||
self.mesh_normalize_scale_factor = scale_factor / scale
|
||
self.mesh_normalize_scale_center = center[None, :]
|
||
else:
|
||
self.scale_factor = 1.0
|
||
self.mesh_normalize_scale_factor = 1.0
|
||
self.mesh_normalize_scale_center = np.array([[0, 0, 0]])
|
||
|
||
self.vtx_pos = mx.array(vp)
|
||
self.pos_idx = mx.array(pi)
|
||
|
||
if vtx_uv is not None and uv_idx is not None:
|
||
uv = np.array(vtx_uv, dtype=np.float32)
|
||
uv[:, 1] = 1.0 - uv[:, 1]
|
||
self.vtx_uv = mx.array(uv)
|
||
self.uv_idx = mx.array(np.array(uv_idx, dtype=np.int32))
|
||
self.extract_textiles()
|
||
else:
|
||
self.vtx_uv = None
|
||
self.uv_idx = None
|
||
|
||
def load_mesh(self, mesh, scale_factor=1.15, auto_center=True):
|
||
"""Load mesh from file path or trimesh object."""
|
||
import trimesh
|
||
|
||
if isinstance(mesh, str):
|
||
mesh = trimesh.load(mesh, process=False, force="mesh")
|
||
|
||
vtx_pos = np.array(mesh.vertices, dtype=np.float32)
|
||
pos_idx = np.array(mesh.faces, dtype=np.int32)
|
||
|
||
vtx_uv, uv_idx = None, None
|
||
if hasattr(mesh.visual, "uv") and mesh.visual.uv is not None:
|
||
vtx_uv = np.array(mesh.visual.uv, dtype=np.float32)
|
||
uv_idx = pos_idx # same topology
|
||
|
||
self.set_mesh(
|
||
vtx_pos, pos_idx, vtx_uv=vtx_uv, uv_idx=uv_idx,
|
||
scale_factor=scale_factor, auto_center=auto_center,
|
||
)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Internal pipeline
|
||
# ------------------------------------------------------------------
|
||
|
||
def _create_view_state(self, elev, azim, camera_distance=None,
|
||
center=None, resolution=None):
|
||
"""Build clip-space vertices and camera-space positions."""
|
||
if camera_distance is None:
|
||
camera_distance = self.camera_distance
|
||
if resolution is None:
|
||
resolution = self.default_resolution
|
||
|
||
mv = get_mv_matrix(elev, azim, camera_distance, center)
|
||
pos_camera = transform_pos(mv, self.vtx_pos, keepdim=True) # (N, 4)
|
||
pos_clip = transform_pos(self.camera_proj_mat, pos_camera) # (1, N, 4)
|
||
return pos_camera, pos_clip, resolution
|
||
|
||
def _rasterize(self, pos_clip, resolution):
|
||
"""Run rasterization, return packed rast_out (1, H, W, 4)."""
|
||
findices, bary = self.raster.rasterize(
|
||
pos_clip, self.pos_idx, resolution
|
||
)
|
||
# Pack [bary(3) | findices(1)] -> (1, H, W, 4)
|
||
rast_out = mx.concatenate(
|
||
[bary, mx.expand_dims(findices.astype(mx.float32), axis=-1)],
|
||
axis=-1,
|
||
)
|
||
return mx.expand_dims(rast_out, axis=0)
|
||
|
||
def _rasterize_tiled(self, pos_clip, resolution, tiles=4):
|
||
"""Tile rasterization across clip-space quadrants.
|
||
|
||
At 4096² the single-dispatch raster exceeds Metal's command-buffer
|
||
budget. Split clip space into ``tiles`` × ``tiles`` regions; for
|
||
each region, scale vertex clip coords so the region maps to
|
||
[-1, 1]² and rasterize at (H/tiles, W/tiles). Concatenate.
|
||
|
||
Assumes ``tiles`` divides H and W evenly. Preserves the row-to-
|
||
clip-y convention of the underlying rasterizer (verified by
|
||
comparing tiled output to single-dispatch output at 2048²).
|
||
"""
|
||
H, W = int(resolution[0]), int(resolution[1])
|
||
if H % tiles != 0 or W % tiles != 0:
|
||
return self._rasterize(pos_clip, resolution)
|
||
tH, tW = H // tiles, W // tiles
|
||
|
||
verts = pos_clip[0] if pos_clip.ndim == 3 else pos_clip
|
||
vx, vy, vz, vw = verts[:, 0], verts[:, 1], verts[:, 2], verts[:, 3]
|
||
|
||
rows_fi, rows_bary = [], []
|
||
for i in range(tiles):
|
||
cols_fi, cols_bary = [], []
|
||
# Tile i covers rows [i*tH, (i+1)*tH). Row 0 = first tile output.
|
||
# Matches the full-raster convention (tiled(0,0) == full[0:tH, 0:tW]).
|
||
y_lo = -1.0 + 2.0 * i / tiles
|
||
y_hi = -1.0 + 2.0 * (i + 1) / tiles
|
||
mid_y = 0.5 * (y_lo + y_hi)
|
||
span_y = 0.5 * (y_hi - y_lo)
|
||
for j in range(tiles):
|
||
x_lo = -1.0 + 2.0 * j / tiles
|
||
x_hi = -1.0 + 2.0 * (j + 1) / tiles
|
||
mid_x = 0.5 * (x_lo + x_hi)
|
||
span_x = 0.5 * (x_hi - x_lo)
|
||
# Homogeneous-safe clip transform: NDC' = (NDC - mid)/span
|
||
# ⇒ clip' = (clip - w*mid)/span (w unchanged).
|
||
new_verts = mx.stack([
|
||
(vx - vw * mid_x) / span_x,
|
||
(vy - vw * mid_y) / span_y,
|
||
vz,
|
||
vw,
|
||
], axis=-1)
|
||
fi, bary = self.raster.rasterize(
|
||
new_verts, self.pos_idx, (tH, tW),
|
||
)
|
||
mx.synchronize()
|
||
cols_fi.append(fi)
|
||
cols_bary.append(bary)
|
||
rows_fi.append(mx.concatenate(cols_fi, axis=1))
|
||
rows_bary.append(mx.concatenate(cols_bary, axis=1))
|
||
findices = mx.concatenate(rows_fi, axis=0)
|
||
bary = mx.concatenate(rows_bary, axis=0)
|
||
rast_out = mx.concatenate(
|
||
[bary, mx.expand_dims(findices.astype(mx.float32), axis=-1)],
|
||
axis=-1,
|
||
)
|
||
return mx.expand_dims(rast_out, axis=0)
|
||
|
||
def _compute_face_normals(self, triangles):
|
||
"""Face normals from (F, 3, 3) triangle vertices."""
|
||
e1 = triangles[:, 1, :] - triangles[:, 0, :]
|
||
e2 = triangles[:, 2, :] - triangles[:, 0, :]
|
||
# Cross product
|
||
n = mx.stack([
|
||
e1[:, 1] * e2[:, 2] - e1[:, 2] * e2[:, 1],
|
||
e1[:, 2] * e2[:, 0] - e1[:, 0] * e2[:, 2],
|
||
e1[:, 0] * e2[:, 1] - e1[:, 1] * e2[:, 0],
|
||
], axis=-1)
|
||
norm = mx.sqrt((n * n).sum(axis=-1, keepdims=True) + 1e-12)
|
||
return n / norm
|
||
|
||
def _get_normals(self, pos_camera, pos_clip, resolution, use_abs_coor=False):
|
||
"""Compute per-pixel normals via face or vertex shading."""
|
||
if use_abs_coor:
|
||
mesh_tris = self.vtx_pos[self.pos_idx] # (F, 3, 3)
|
||
else:
|
||
pc3 = pos_camera[:, :3] / pos_camera[:, 3:4]
|
||
mesh_tris = pc3[self.pos_idx]
|
||
|
||
face_normals = self._compute_face_normals(mesh_tris)
|
||
rast_out = self._rasterize(pos_clip, resolution)
|
||
|
||
if self.shader_type == "vertex":
|
||
vertex_normals = self._vertex_normals_from_faces(face_normals)
|
||
normal = self.raster.interpolate(
|
||
mx.expand_dims(vertex_normals, 0),
|
||
rast_out[0, ..., -1].astype(mx.int32),
|
||
rast_out[0, ..., :-1],
|
||
self.pos_idx,
|
||
)
|
||
else:
|
||
# Face shader: direct face-normal lookup per pixel
|
||
tri_ids = rast_out[..., 3].astype(mx.int32) # (1, H, W)
|
||
mask = (tri_ids > 0).astype(mx.float32)
|
||
tri_ids_0 = mx.maximum(tri_ids - 1, 0) # 0-indexed
|
||
normal = face_normals[tri_ids_0[0]] # (H, W, 3)
|
||
normal = normal * mask[0, ..., None]
|
||
normal = mx.expand_dims(normal, 0) # (1, H, W, 3)
|
||
|
||
return normal, rast_out
|
||
|
||
def _vertex_normals_from_faces(self, face_normals):
|
||
"""Compute vertex normals as average of adjacent face normals."""
|
||
num_verts = self.vtx_pos.shape[0]
|
||
fi = np.array(self.pos_idx)
|
||
fn = np.array(face_normals)
|
||
vn = np.zeros((num_verts, 3), dtype=np.float32)
|
||
counts = np.zeros((num_verts, 1), dtype=np.float32)
|
||
for c in range(3):
|
||
np.add.at(vn, fi[:, c], fn)
|
||
np.add.at(counts, fi[:, c], 1.0)
|
||
counts = np.maximum(counts, 1.0)
|
||
vn = vn / counts
|
||
norms = np.maximum(np.linalg.norm(vn, axis=-1, keepdims=True), 1e-12)
|
||
return mx.array(vn / norms)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Public render API
|
||
# ------------------------------------------------------------------
|
||
|
||
@staticmethod
|
||
def _format_output(image_np, return_type="np"):
|
||
"""Convert numpy output to requested format."""
|
||
if return_type == "pl":
|
||
from PIL import Image as PILImage
|
||
img = np.clip(image_np, 0, 1)
|
||
return PILImage.fromarray((img * 255).astype(np.uint8))
|
||
return image_np
|
||
|
||
def render_normal(self, elev, azim, resolution=None, bg_color=(1, 1, 1),
|
||
use_abs_coor=False, camera_distance=None, center=None,
|
||
return_type="np"):
|
||
"""Render surface normals from a given viewpoint.
|
||
|
||
Returns:
|
||
(H, W, 3) numpy array (return_type="np") or PIL Image ("pl").
|
||
"""
|
||
pos_camera, pos_clip, res = self._create_view_state(
|
||
elev, azim, camera_distance, center, resolution,
|
||
)
|
||
normal, rast_out = self._get_normals(
|
||
pos_camera, pos_clip, res, use_abs_coor,
|
||
)
|
||
visible = mx.clip(rast_out[..., -1:], 0, 1)
|
||
bg = mx.array(list(bg_color), dtype=mx.float32).reshape(1, 1, 1, 3)
|
||
result = normal * visible + bg * (1.0 - visible)
|
||
result = (result + 1.0) * 0.5
|
||
mx.synchronize()
|
||
return self._format_output(np.array(result[0]), return_type)
|
||
|
||
def render_position(self, elev, azim, resolution=None, bg_color=(1, 1, 1),
|
||
camera_distance=None, center=None, return_type="np"):
|
||
"""Render world-space positions from a given viewpoint.
|
||
|
||
Returns:
|
||
(H, W, 3) numpy array (return_type="np") or PIL Image ("pl").
|
||
"""
|
||
_, pos_clip, res = self._create_view_state(
|
||
elev, azim, camera_distance, center, resolution,
|
||
)
|
||
rast_out = self._rasterize(pos_clip, res)
|
||
findices = rast_out[0, ..., -1].astype(mx.int32)
|
||
bary = rast_out[0, ..., :-1]
|
||
|
||
tex_pos = 0.5 - self.vtx_pos[:, :3] / self.scale_factor
|
||
position = self.raster.interpolate(
|
||
mx.expand_dims(tex_pos, 0), findices, bary, self.pos_idx,
|
||
)
|
||
|
||
visible = mx.clip(rast_out[..., -1:], 0, 1)
|
||
bg = mx.array(list(bg_color), dtype=mx.float32).reshape(1, 1, 1, 3)
|
||
result = position * visible + bg * (1.0 - visible)
|
||
mx.synchronize()
|
||
return self._format_output(np.array(result[0]), return_type)
|
||
|
||
def render_alpha(self, elev, azim, resolution=None,
|
||
camera_distance=None, center=None, return_type="np"):
|
||
"""Render face-index map from a given viewpoint.
|
||
|
||
Returns face indices (1-indexed, 0=background) matching the original
|
||
MeshRender behavior. Shape: (1, H, W, 1) for return_type="np".
|
||
"""
|
||
_, pos_clip, res = self._create_view_state(
|
||
elev, azim, camera_distance, center, resolution,
|
||
)
|
||
rast_out = self._rasterize(pos_clip, res)
|
||
# Return face indices (int), same as original: rast_out[..., -1:].long()
|
||
face_ids = np.array(rast_out[..., -1:])[0].astype(np.int64)
|
||
# Shape (1, H, W, 1) to match original pipeline expectations
|
||
return face_ids[None, ...]
|
||
|
||
# ------------------------------------------------------------------
|
||
# Mesh & texture management
|
||
# ------------------------------------------------------------------
|
||
|
||
def set_default_render_resolution(self, default_resolution):
|
||
if isinstance(default_resolution, int):
|
||
default_resolution = (default_resolution, default_resolution)
|
||
self.default_resolution = default_resolution
|
||
|
||
def set_boundary_unreliable_scale(self, scale):
|
||
self.bake_unreliable_kernel_size = int(
|
||
(scale / 512)
|
||
* max(self.default_resolution[0], self.default_resolution[1])
|
||
)
|
||
|
||
def get_face_areas(self, from_one_index=False):
|
||
"""Compute area of each triangle face.
|
||
|
||
Returns:
|
||
numpy array of face areas. If from_one_index, prepends a zero
|
||
so areas[face_id] works with 1-indexed face IDs.
|
||
"""
|
||
vp = np.array(self.vtx_pos)
|
||
fi = np.array(self.pos_idx)
|
||
v0 = vp[fi[:, 0]]
|
||
v1 = vp[fi[:, 1]]
|
||
v2 = vp[fi[:, 2]]
|
||
cross = np.cross(v1 - v0, v2 - v0)
|
||
areas = np.linalg.norm(cross, axis=-1) * 0.5
|
||
if from_one_index:
|
||
areas = np.insert(areas, 0, 0.0)
|
||
return areas
|
||
|
||
def set_texture(self, tex, force_set=False):
|
||
"""Store diffuse texture (numpy uint8 or float, or PIL Image)."""
|
||
from PIL import Image as PILImage
|
||
if isinstance(tex, PILImage.Image):
|
||
tex = np.array(tex).astype(np.float32) / 255.0
|
||
elif isinstance(tex, np.ndarray) and tex.dtype == np.uint8:
|
||
tex = tex.astype(np.float32) / 255.0
|
||
self.tex = tex
|
||
|
||
def get_texture(self):
|
||
"""Return diffuse texture as numpy float32 array."""
|
||
return self.tex if hasattr(self, "tex") and self.tex is not None else None
|
||
|
||
def set_texture_mr(self, mr, force_set=False):
|
||
"""Store metallic-roughness texture."""
|
||
from PIL import Image as PILImage
|
||
if isinstance(mr, PILImage.Image):
|
||
mr = np.array(mr).astype(np.float32) / 255.0
|
||
elif isinstance(mr, np.ndarray) and mr.dtype == np.uint8:
|
||
mr = mr.astype(np.float32) / 255.0
|
||
self.tex_mr = mr
|
||
|
||
def get_texture_mr(self):
|
||
metallic, roughness = None, None
|
||
if hasattr(self, "tex_mr") and self.tex_mr is not None:
|
||
mr = self.tex_mr
|
||
metallic = np.repeat(mr[:, :, 0:1], 3, axis=2)
|
||
roughness = np.repeat(mr[:, :, 1:2], 3, axis=2)
|
||
return metallic, roughness
|
||
|
||
def get_texture_normal(self):
|
||
if hasattr(self, "tex_normalMap") and self.tex_normalMap is not None:
|
||
return self.tex_normalMap
|
||
return None
|
||
|
||
def get_mesh(self, normalize=True):
|
||
"""Return mesh geometry as numpy arrays.
|
||
|
||
Applies inverse coordinate transform to restore original space.
|
||
"""
|
||
vtx_pos = np.array(self.vtx_pos)
|
||
pos_idx = np.array(self.pos_idx)
|
||
|
||
if not normalize:
|
||
vtx_pos = vtx_pos / self.mesh_normalize_scale_factor
|
||
vtx_pos = vtx_pos + self.mesh_normalize_scale_center
|
||
|
||
# Inverse of set_mesh transform: undo swap Y/Z, then undo flip X/Y
|
||
vtx_pos[:, [1, 2]] = vtx_pos[:, [2, 1]].copy()
|
||
vtx_pos[:, [0, 1]] = -vtx_pos[:, [0, 1]]
|
||
|
||
vtx_uv, uv_idx = None, None
|
||
if self.vtx_uv is not None:
|
||
vtx_uv = np.array(self.vtx_uv)
|
||
vtx_uv[:, 1] = 1.0 - vtx_uv[:, 1]
|
||
uv_idx = np.array(self.uv_idx)
|
||
|
||
return vtx_pos, pos_idx, vtx_uv, uv_idx
|
||
|
||
def save_mesh(self, mesh_path, downsample=False):
|
||
"""Save textured mesh to OBJ file."""
|
||
import cv2
|
||
import trimesh
|
||
|
||
vtx_pos, pos_idx, vtx_uv, uv_idx = self.get_mesh(normalize=False)
|
||
texture = self.get_texture()
|
||
|
||
if texture is not None and downsample:
|
||
h, w = texture.shape[0] // 2, texture.shape[1] // 2
|
||
texture = cv2.resize(texture, (w, h))
|
||
|
||
mesh = trimesh.Trimesh(
|
||
vertices=vtx_pos,
|
||
faces=pos_idx,
|
||
process=False,
|
||
)
|
||
if vtx_uv is not None and texture is not None:
|
||
from PIL import Image as PILImage
|
||
tex_img = PILImage.fromarray(
|
||
(np.clip(texture, 0, 1) * 255).astype(np.uint8)
|
||
)
|
||
# Use PBRMaterial so GLB exports tag the baseColorTexture as
|
||
# sRGB. SimpleMaterial doesn't carry color-space metadata and
|
||
# GLB viewers assumed the texture was linear, applying an
|
||
# unwanted gamma curve that darkened / desaturated the mesh
|
||
# relative to the atlas PNG.
|
||
mr_tex_img = None
|
||
mr_atlas = getattr(self, "tex_mr", None)
|
||
if mr_atlas is not None:
|
||
# Repack for glTF: the diffusion MR atlas uses R=metallic,
|
||
# G=roughness (PT convention, see MeshRender.get_texture_mr),
|
||
# but glTF metallicRoughnessTexture expects B=metallic,
|
||
# G=roughness, R ignored (occlusion in some pipelines).
|
||
# Without this swap, glTF viewers read metallic=0 and the
|
||
# mesh looks uniformly dielectric.
|
||
mr_u8 = (np.clip(mr_atlas, 0, 1) * 255).astype(np.uint8)
|
||
if mr_u8.ndim == 2:
|
||
mr_u8 = np.stack([mr_u8, mr_u8, mr_u8], axis=-1)
|
||
metallic_ch = mr_u8[..., 0]
|
||
roughness_ch = mr_u8[..., 1]
|
||
mr_gltf = np.stack(
|
||
[np.zeros_like(metallic_ch), roughness_ch, metallic_ch],
|
||
axis=-1,
|
||
)
|
||
mr_tex_img = PILImage.fromarray(mr_gltf)
|
||
|
||
kwargs = dict(
|
||
name="paint_pbr",
|
||
baseColorTexture=tex_img,
|
||
)
|
||
if mr_tex_img is not None:
|
||
kwargs["metallicRoughnessTexture"] = mr_tex_img
|
||
kwargs["metallicFactor"] = 1.0
|
||
kwargs["roughnessFactor"] = 1.0
|
||
else:
|
||
kwargs["metallicFactor"] = 0.0
|
||
kwargs["roughnessFactor"] = 1.0
|
||
material = trimesh.visual.material.PBRMaterial(**kwargs)
|
||
# doubleSided stops backface culling that was making the
|
||
# silhouettes see-through.
|
||
material.doubleSided = True
|
||
visuals = trimesh.visual.TextureVisuals(
|
||
uv=vtx_uv, material=material,
|
||
)
|
||
mesh.visual = visuals
|
||
|
||
mesh.export(mesh_path)
|
||
# Expose the in-memory trimesh so callers can export GLB directly.
|
||
# Going OBJ -> trimesh.load -> GLB would drop the PBRMaterial
|
||
# (OBJ/MTL can't carry metallicRoughnessTexture, doubleSided, or
|
||
# sRGB hints) and silently erase our PBR work.
|
||
self._last_exported_mesh = mesh
|
||
|
||
# ------------------------------------------------------------------
|
||
# Texture-space operations
|
||
# ------------------------------------------------------------------
|
||
|
||
def extract_textiles(self):
|
||
"""Rasterize mesh in UV space to build texture-space geometry maps.
|
||
|
||
Populates self.tex_position, tex_normal, tex_grid, texture_indices.
|
||
Required by back_project(method="back_sample").
|
||
"""
|
||
if self.vtx_uv is None or self.uv_idx is None:
|
||
return
|
||
|
||
# Build clip-space coords from UVs: (u, v, 0, 1) * 2 - 1
|
||
vnum = self.vtx_uv.shape[0]
|
||
vtx_uv_clip = mx.concatenate([
|
||
self.vtx_uv,
|
||
mx.zeros((vnum, 1)),
|
||
mx.ones((vnum, 1)),
|
||
], axis=1) * 2.0 - 1.0
|
||
vtx_uv_clip = mx.expand_dims(vtx_uv_clip, 0) # (1, V, 4)
|
||
|
||
# Rasterize in UV space. At ≥4096² the single-dispatch raster
|
||
# exceeds Metal's command-buffer budget, so tile it.
|
||
th, tw = self.texture_size
|
||
max_dim = max(th, tw)
|
||
tiles = 1
|
||
if max_dim >= 4096:
|
||
tiles = 4
|
||
elif max_dim >= 3072:
|
||
tiles = 2
|
||
if tiles > 1:
|
||
rast_out = self._rasterize_tiled(vtx_uv_clip, self.texture_size, tiles=tiles)
|
||
else:
|
||
rast_out = self._rasterize(vtx_uv_clip, self.texture_size)
|
||
fi = rast_out[0, ..., -1].astype(mx.int32)
|
||
bary = rast_out[0, ..., :-1]
|
||
|
||
# Interpolate world positions in UV space
|
||
position = self.raster.interpolate(
|
||
mx.expand_dims(self.vtx_pos, 0), fi, bary, self.pos_idx,
|
||
)[0] # (H, W, 3)
|
||
|
||
# Compute face normals and per-pixel normals (face shader)
|
||
face_normals = self._compute_face_normals(self.vtx_pos[self.pos_idx])
|
||
tri_ids = rast_out[0, ..., 3].astype(mx.int32)
|
||
mask = tri_ids > 0
|
||
tri_ids_0 = mx.maximum(tri_ids - 1, 0)
|
||
position_normal = face_normals[tri_ids_0] # (H, W, 3)
|
||
position_normal = position_normal * mask[..., None].astype(mx.float32)
|
||
|
||
visible_mask = mx.clip(rast_out[0, ..., -1], 0, 1) # (H, W)
|
||
mx.synchronize()
|
||
|
||
# Extract visible pixels to flat arrays (numpy for indexing)
|
||
vis_np = np.array(visible_mask).reshape(-1)
|
||
mask_flat = vis_np > 0
|
||
|
||
pos_np = np.array(position).reshape(-1, 3)
|
||
norm_np = np.array(position_normal).reshape(-1, 3)
|
||
|
||
th, tw = self.texture_size
|
||
row, col = np.meshgrid(np.arange(th), np.arange(tw), indexing="ij")
|
||
grid_np = np.stack([row, col], axis=-1).reshape(-1, 2)
|
||
|
||
pos_visible = pos_np[mask_flat]
|
||
norm_visible = norm_np[mask_flat]
|
||
grid_visible = grid_np[mask_flat]
|
||
|
||
# Add homogeneous w=1
|
||
pos_visible = np.concatenate(
|
||
[pos_visible, np.ones((pos_visible.shape[0], 1), dtype=np.float32)],
|
||
axis=1,
|
||
)
|
||
|
||
# Build reverse index: texture pixel -> flat visible index
|
||
texture_indices = np.full(th * tw, -1, dtype=np.int64)
|
||
flat_idx = grid_visible[:, 0] * tw + grid_visible[:, 1]
|
||
texture_indices[flat_idx] = np.arange(len(grid_visible))
|
||
|
||
self.tex_position = pos_visible # (K, 4) float32
|
||
self.tex_normal = norm_visible # (K, 3) float32
|
||
self.tex_grid = grid_visible # (K, 2) int
|
||
self.texture_indices = texture_indices.reshape(th, tw) # (H, W)
|
||
|
||
# Per-texel face id (1-indexed, 0 = empty), for per-face bake merging
|
||
self.tex_face_id = np.array(tri_ids).astype(np.int32) # (H, W)
|
||
|
||
def uv_feature_map(self, vert_feat):
|
||
"""Map per-vertex features to UV texture space.
|
||
|
||
Args:
|
||
vert_feat: (N, C) mx.array of per-vertex features.
|
||
|
||
Returns:
|
||
(H, W, C) mx.array feature map.
|
||
"""
|
||
vtx_uv_clip = mx.concatenate([
|
||
self.vtx_uv * 2 - 1,
|
||
mx.zeros((self.vtx_uv.shape[0], 1)),
|
||
mx.ones((self.vtx_uv.shape[0], 1)),
|
||
], axis=1)
|
||
vtx_uv_clip = mx.expand_dims(vtx_uv_clip, 0)
|
||
|
||
rast_out = self._rasterize(vtx_uv_clip, self.texture_size)
|
||
fi = rast_out[0, ..., -1].astype(mx.int32)
|
||
bary = rast_out[0, ..., :-1]
|
||
|
||
feat_map = self.raster.interpolate(
|
||
mx.expand_dims(vert_feat, 0), fi, bary, self.uv_idx,
|
||
)
|
||
return feat_map[0] # (H, W, C)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Texture baking
|
||
# ------------------------------------------------------------------
|
||
|
||
@staticmethod
|
||
def _render_sketch_from_depth(depth_image_np):
|
||
"""Edge detection on depth map (CPU, cv2)."""
|
||
import cv2
|
||
depth_u8 = (depth_image_np * 255).astype(np.uint8)
|
||
edges = cv2.Canny(depth_u8, 30, 80)
|
||
return (edges.astype(np.float32) / 255.0)[..., None] # (H, W, 1)
|
||
|
||
@staticmethod
|
||
def _erode_boundary(mask_np, sketch_np, kernel_size):
|
||
"""Erode visible mask near silhouette boundaries (CPU, cv2)."""
|
||
import cv2
|
||
if kernel_size <= 0:
|
||
return mask_np
|
||
ks = kernel_size * 2 + 1
|
||
kernel = np.ones((ks, ks), dtype=np.float32)
|
||
|
||
# Erode visible mask: pixels near background become unreliable
|
||
inv = 1.0 - mask_np[..., 0]
|
||
eroded = cv2.filter2D(inv, -1, kernel)
|
||
mask_np = mask_np * (eroded <= 0).astype(np.float32)[..., None]
|
||
|
||
# Also exclude pixels near depth edges
|
||
sketch_dilated = cv2.filter2D(sketch_np[..., 0], -1, kernel)
|
||
mask_np = mask_np * (sketch_dilated < 0.5).astype(np.float32)[..., None]
|
||
return mask_np
|
||
|
||
def back_project(self, image, elev, azim, camera_distance=None,
|
||
center=None):
|
||
"""Back-project a rendered image onto UV texture space.
|
||
|
||
Uses the "back_sample" method: projects texture-space positions
|
||
into image space and bilinearly samples colors.
|
||
|
||
Args:
|
||
image: (H, W, C) numpy array or PIL Image, values in [0, 1].
|
||
elev: Camera elevation in degrees.
|
||
azim: Camera azimuth in degrees.
|
||
|
||
Returns:
|
||
texture: (H, W, C) numpy array in UV space.
|
||
cos_map: (H, W, 1) numpy array cosine weights.
|
||
boundary_map: (H, W, 1) numpy array boundary mask.
|
||
"""
|
||
from PIL import Image as PILImage
|
||
|
||
if isinstance(image, PILImage.Image):
|
||
image = np.array(image).astype(np.float32) / 255.0
|
||
elif not isinstance(image, np.ndarray):
|
||
image = np.array(image)
|
||
if image.ndim == 2:
|
||
image = image[..., None]
|
||
image = image.astype(np.float32)
|
||
|
||
resolution = image.shape[:2] # (H, W)
|
||
channel = image.shape[2]
|
||
th, tw = self.texture_size
|
||
|
||
if camera_distance is None:
|
||
camera_distance = self.camera_distance
|
||
|
||
# --- Camera setup ---
|
||
proj = self.camera_proj_mat
|
||
mv = get_mv_matrix(elev, azim, camera_distance, center)
|
||
pos_camera = transform_pos(mv, self.vtx_pos, keepdim=True) # (N, 4)
|
||
pos_clip = transform_pos(self.camera_proj_mat, pos_camera) # (1, N, 4)
|
||
|
||
pc3 = pos_camera[:, :3] / pos_camera[:, 3:4]
|
||
|
||
# --- Face normals in camera space ---
|
||
mesh_tris = pc3[self.pos_idx] # (F, 3, 3)
|
||
face_normals = self._compute_face_normals(mesh_tris)
|
||
|
||
# --- Rasterize from camera view ---
|
||
rast_out = self._rasterize(pos_clip, resolution)
|
||
fi = rast_out[0, ..., -1].astype(mx.int32)
|
||
bary = rast_out[0, ..., :-1]
|
||
|
||
visible = mx.clip(rast_out[0, ..., -1:], 0, 1) # (H, W, 1)
|
||
|
||
# Per-pixel normals (face shader)
|
||
tri_ids = fi
|
||
tri_mask = (tri_ids > 0).astype(mx.float32)
|
||
tri_ids_0 = mx.maximum(tri_ids - 1, 0)
|
||
normal = face_normals[tri_ids_0] * tri_mask[..., None] # (H, W, 3)
|
||
|
||
# Depth
|
||
depth_attr = pc3[:, 2:3].reshape(1, -1, 1)
|
||
depth = self.raster.interpolate(depth_attr, fi, bary, self.pos_idx)
|
||
depth = depth[0] # (H, W, 1)
|
||
mx.synchronize()
|
||
|
||
visible_np = np.array(visible)
|
||
normal_np = np.array(normal)
|
||
depth_np = np.array(depth)
|
||
|
||
# Depth normalization for sketch
|
||
vis_flat = visible_np.reshape(-1) > 0
|
||
if vis_flat.sum() == 0:
|
||
return (np.zeros((*self.texture_size, channel), dtype=np.float32),
|
||
np.zeros((*self.texture_size, 1), dtype=np.float32),
|
||
np.zeros((*self.texture_size, 1), dtype=np.float32))
|
||
d_vis = depth_np.reshape(-1)[vis_flat]
|
||
d_max, d_min = d_vis.max(), d_vis.min()
|
||
depth_norm = (depth_np - d_min) / max(d_max - d_min, 1e-8)
|
||
depth_image = depth_norm * visible_np
|
||
|
||
sketch_np = self._render_sketch_from_depth(depth_image[..., 0])
|
||
|
||
# Cosine weighting (view angle vs surface normal)
|
||
lookat = np.array([0, 0, -1], dtype=np.float32)
|
||
normal_flat = normal_np.reshape(-1, 3)
|
||
cos_flat = (normal_flat * lookat).sum(axis=-1)
|
||
cos_np = cos_flat.reshape(resolution[0], resolution[1], 1)
|
||
cos_thres = np.cos(np.radians(self.bake_angle_thres))
|
||
cos_np[cos_np < cos_thres] = 0.0
|
||
|
||
# Boundary erosion
|
||
visible_np = self._erode_boundary(
|
||
visible_np, sketch_np, self.bake_unreliable_kernel_size,
|
||
)
|
||
cos_np[visible_np[..., 0] == 0] = 0.0
|
||
|
||
# --- Back-sample: project tex_position to image space ---
|
||
if self.tex_position is None:
|
||
raise RuntimeError("extract_textiles() must be called first")
|
||
|
||
img_proj = np.array([
|
||
[proj[0, 0], 0, 0, 0],
|
||
[0, proj[1, 1], 0, 0],
|
||
[0, 0, 1, 0],
|
||
[0, 0, 0, 1],
|
||
], dtype=np.float32)
|
||
|
||
mv_f = mv.astype(np.float32)
|
||
tex_pos = self.tex_position # (K, 4)
|
||
v_proj = tex_pos @ mv_f.T @ img_proj # (K, 4)
|
||
|
||
# Check which tex pixels project inside the image
|
||
inner = (
|
||
(v_proj[:, 0] >= -1.0) & (v_proj[:, 0] <= 1.0)
|
||
& (v_proj[:, 1] >= -1.0) & (v_proj[:, 1] <= 1.0)
|
||
)
|
||
inner_idx = np.where(inner)[0]
|
||
|
||
img_x = np.clip(
|
||
((np.clip(v_proj[:, 0], -1, 1) * 0.5 + 0.5) * resolution[0]).astype(np.int64),
|
||
0, resolution[0] - 1,
|
||
)
|
||
img_y = np.clip(
|
||
((np.clip(v_proj[:, 1], -1, 1) * 0.5 + 0.5) * resolution[1]).astype(np.int64),
|
||
0, resolution[1] - 1,
|
||
)
|
||
|
||
indices = img_y * resolution[0] + img_x
|
||
depth_flat = depth_np.reshape(-1)
|
||
vis_mask_flat = visible_np.reshape(-1)
|
||
cos_flat = cos_np.reshape(-1)
|
||
|
||
sampled_z = depth_flat[indices]
|
||
sampled_m = vis_mask_flat[indices]
|
||
sampled_w = cos_flat[indices]
|
||
v_z = v_proj[:, 2]
|
||
|
||
depth_thres = 3e-3
|
||
valid_mask = (np.abs(v_z - sampled_z) < depth_thres) & (sampled_m * sampled_w > 0)
|
||
valid_idx = np.where(valid_mask)[0]
|
||
valid_idx = np.intersect1d(valid_idx, inner_idx)
|
||
|
||
# Bilinear sampling from image
|
||
wx = ((v_proj[:, 0] * 0.5 + 0.5) * resolution[0] - img_x)[valid_idx].reshape(-1, 1)
|
||
wy = ((v_proj[:, 1] * 0.5 + 0.5) * resolution[1] - img_y)[valid_idx].reshape(-1, 1)
|
||
ix = img_x[valid_idx]
|
||
iy = img_y[valid_idx]
|
||
ix_r = np.clip(ix + 1, 0, resolution[0] - 1)
|
||
iy_r = np.clip(iy + 1, 0, resolution[1] - 1)
|
||
|
||
rgb = image.reshape(-1, channel)
|
||
i00 = iy * resolution[0] + ix
|
||
i10 = iy * resolution[0] + ix_r
|
||
i01 = iy_r * resolution[0] + ix
|
||
i11 = iy_r * resolution[0] + ix_r
|
||
sampled_rgb = (
|
||
(rgb[i00] * (1 - wx) + rgb[i10] * wx) * (1 - wy)
|
||
+ (rgb[i01] * (1 - wx) + rgb[i11] * wx) * wy
|
||
)
|
||
|
||
sampled_b = sketch_np.reshape(-1)[indices[valid_idx]]
|
||
sampled_w_valid = sampled_w[valid_idx]
|
||
|
||
# Write to texture space
|
||
texture = np.zeros((th * tw, channel), dtype=np.float32)
|
||
cos_map = np.zeros(th * tw, dtype=np.float32)
|
||
boundary_map = np.zeros(th * tw, dtype=np.float32)
|
||
|
||
tex_grid = self.tex_grid[valid_idx]
|
||
tex_flat_idx = tex_grid[:, 0] * tw + tex_grid[:, 1]
|
||
texture[tex_flat_idx] = sampled_rgb
|
||
cos_map[tex_flat_idx] = sampled_w_valid
|
||
boundary_map[tex_flat_idx] = sampled_b
|
||
|
||
return (
|
||
texture.reshape(th, tw, channel),
|
||
cos_map.reshape(th, tw, 1),
|
||
boundary_map.reshape(th, tw, 1),
|
||
)
|
||
|
||
def fast_bake_texture(self, textures, cos_maps, mode="weighted"):
|
||
"""Merge multiple view textures into a single UV atlas.
|
||
|
||
Three modes (PT parity is default):
|
||
- ``"weighted"`` (default, PyTorch parity): cosine-weighted
|
||
average of all views. Smooth across face boundaries; matches
|
||
MeshRender.fast_bake_texture exactly.
|
||
- ``"face_wta"``: pick one view per UV face. Sharper but creates
|
||
visible discontinuities at mesh-face seams where adjacent
|
||
faces pick different winning views.
|
||
- ``"wta"``: per-texel argmax. Fastest, stripiest.
|
||
|
||
Args:
|
||
textures: list of (H, W, C) numpy arrays.
|
||
cos_maps: list of (H, W, 1) numpy arrays (already weighted).
|
||
mode: "wta" or "weighted".
|
||
|
||
Returns:
|
||
texture_merge: (H, W, C) numpy array.
|
||
trust_map: (H, W, 1) boolean numpy array.
|
||
"""
|
||
if not textures:
|
||
raise ValueError("fast_bake_texture: no textures provided")
|
||
|
||
if mode == "face_wta":
|
||
# Per-face WTA: aggregate cos per (face, view), pick best view per
|
||
# face, then bake all of that face's texels from that view.
|
||
if getattr(self, "tex_face_id", None) is None:
|
||
# Fall back to per-texel WTA if face ids aren't available
|
||
mode = "wta"
|
||
else:
|
||
stack_t = np.stack(textures, axis=0) # (V, H, W, C)
|
||
stack_c = np.stack([c[..., 0] for c in cos_maps], axis=0) # (V, H, W)
|
||
face_ids = self.tex_face_id # (H, W) 1-indexed
|
||
V = stack_c.shape[0]
|
||
n_faces = int(face_ids.max())
|
||
# Sum cos per (view, face)
|
||
flat_face = face_ids.reshape(-1)
|
||
valid = flat_face > 0
|
||
face_idx = (flat_face[valid] - 1).astype(np.int64)
|
||
stack_c_flat = stack_c.reshape(V, -1)[:, valid] # (V, K)
|
||
cos_per_face = np.zeros((V, n_faces), dtype=np.float32)
|
||
for v in range(V):
|
||
np.add.at(cos_per_face[v], face_idx, stack_c_flat[v])
|
||
best_view_per_face = cos_per_face.argmax(axis=0) # (n_faces,)
|
||
# Per-texel best view from face lookup
|
||
best_per_texel = np.zeros_like(face_ids, dtype=np.int64)
|
||
best_per_texel[face_ids > 0] = best_view_per_face[face_ids[face_ids > 0] - 1]
|
||
# Painted = any view contributed at this texel
|
||
best_c = stack_c.max(axis=0)
|
||
H, W = face_ids.shape
|
||
ii, jj = np.meshgrid(np.arange(H), np.arange(W), indexing="ij")
|
||
merged = stack_t[best_per_texel, ii, jj]
|
||
merged[best_c == 0] = 0
|
||
return merged, (best_c > 0)[..., None]
|
||
|
||
if mode == "wta":
|
||
stack_t = np.stack(textures, axis=0)
|
||
stack_c = np.stack([c[..., 0] for c in cos_maps], axis=0)
|
||
best = stack_c.argmax(axis=0)
|
||
best_c = stack_c.max(axis=0)
|
||
H, W = best.shape
|
||
ii, jj = np.meshgrid(np.arange(H), np.arange(W), indexing="ij")
|
||
merged = stack_t[best, ii, jj]
|
||
merged[best_c == 0] = 0
|
||
return merged, (best_c > 0)[..., None]
|
||
|
||
# PyTorch-parity weighted average path
|
||
channel = textures[0].shape[-1]
|
||
th, tw = self.texture_size
|
||
texture_merge = np.zeros((th, tw, channel), dtype=np.float32)
|
||
trust_map = np.zeros((th, tw, 1), dtype=np.float32)
|
||
for texture, cos_map in zip(textures, cos_maps):
|
||
view_sum = (cos_map > 0).sum()
|
||
if view_sum == 0:
|
||
continue
|
||
painted_sum = ((cos_map > 0) & (trust_map > 0)).sum()
|
||
if painted_sum / view_sum > 0.99:
|
||
continue
|
||
texture_merge += texture * cos_map
|
||
trust_map += cos_map
|
||
texture_merge = texture_merge / np.maximum(trust_map, 1e-8)
|
||
return texture_merge, trust_map > 1e-8
|
||
|
||
def bake_texture(self, colors, elevs, azims, camera_distance=None,
|
||
center=None, exp=6, weights=None):
|
||
"""Bake multiple views into a single UV texture.
|
||
|
||
Args:
|
||
colors: list of images (numpy or PIL).
|
||
elevs: list of elevation angles.
|
||
azims: list of azimuth angles.
|
||
exp: exponent for cosine weighting.
|
||
weights: optional per-view weights.
|
||
|
||
Returns:
|
||
texture_merge: (H, W, C) numpy array.
|
||
trust_map: (H, W, 1) boolean numpy array.
|
||
"""
|
||
from PIL import Image as PILImage
|
||
|
||
if weights is None:
|
||
weights = [1.0] * len(colors)
|
||
|
||
textures = []
|
||
cos_maps = []
|
||
for color, elev, azim, weight in zip(colors, elevs, azims, weights):
|
||
if isinstance(color, PILImage.Image):
|
||
color = np.array(color).astype(np.float32) / 255.0
|
||
texture, cos_map, _ = self.back_project(
|
||
color, elev, azim, camera_distance, center,
|
||
)
|
||
cos_map = weight * (cos_map ** exp)
|
||
textures.append(texture)
|
||
cos_maps.append(cos_map)
|
||
|
||
return self.fast_bake_texture(textures, cos_maps)
|
||
|
||
def uv_inpaint(self, texture, mask, method="NS", vertex_inpaint=True):
|
||
"""Inpaint missing regions in UV texture.
|
||
|
||
Two-pass pipeline matching PyTorch's bake flow:
|
||
1. Mesh-aware vertex-color propagation + face-bary rasterization
|
||
(fills most of the UV islands using 3D adjacency).
|
||
2. Gutter fill for the remaining empty texels.
|
||
|
||
For step 2 the default ``"edt"`` mode copies each gutter texel from
|
||
its nearest painted texel via ``scipy.ndimage.distance_transform_edt``.
|
||
It's deterministic and never blends colors across distant UV islands —
|
||
unlike ``"NS"`` (cv2.INPAINT_NS) which diffuses and tends to mix
|
||
unrelated islands together when they happen to be 2D-close in atlas
|
||
space.
|
||
|
||
Args:
|
||
texture: (H, W, C) numpy array in [0, 1].
|
||
mask: (H, W) or (H, W, 1) numpy uint8, 255=keep, 0=inpaint.
|
||
method: "edt" (default, EDT nearest-fill) or "NS" (cv2 NS).
|
||
vertex_inpaint: Run the mesh-aware propagation pass first.
|
||
|
||
Returns:
|
||
(H, W, C) numpy uint8 array.
|
||
"""
|
||
if isinstance(texture, mx.array):
|
||
texture = np.array(texture)
|
||
|
||
if mask.ndim == 3:
|
||
mask = mask[..., 0]
|
||
if mask.dtype != np.uint8:
|
||
mask = (mask * 255).astype(np.uint8)
|
||
|
||
tex_f = np.clip(texture, 0, 1).astype(np.float32)
|
||
|
||
if vertex_inpaint and self.vtx_uv is not None and self.uv_idx is not None:
|
||
from .mesh_inpaint_py import mesh_vertex_inpaint
|
||
tex_f, mask = mesh_vertex_inpaint(
|
||
tex_f, mask,
|
||
np.asarray(self.vtx_pos), np.asarray(self.vtx_uv),
|
||
np.asarray(self.pos_idx), np.asarray(self.uv_idx),
|
||
)
|
||
|
||
if method == "edt":
|
||
from scipy.ndimage import distance_transform_edt
|
||
painted = mask > 0
|
||
if painted.any():
|
||
idx = distance_transform_edt(
|
||
~painted, return_distances=False, return_indices=True,
|
||
)
|
||
tex_f = tex_f[idx[0], idx[1]]
|
||
return (np.clip(tex_f, 0, 1) * 255).astype(np.uint8)
|
||
|
||
# cv2.INPAINT_NS for near-field fill (PT parity), then an explicit
|
||
# EDT edge-padding pass for UV gutter. NS's 3-px radius is smaller
|
||
# than typical UV-island spacing in a 2048^2 atlas, so 3D viewers
|
||
# doing bilinear texture sampling across island boundaries read
|
||
# the NS-smoothed (near-black) gutter and display thin wireframe
|
||
# lines on the mesh. EDT pads every unpainted texel with the
|
||
# nearest painted color — this is the standard post-bake step in
|
||
# production UV pipelines and never affects texels that are
|
||
# visible in 3D (painted).
|
||
import cv2
|
||
from scipy.ndimage import distance_transform_edt
|
||
texture_u8 = (np.clip(tex_f, 0, 1) * 255).astype(np.uint8)
|
||
inpainted = cv2.inpaint(texture_u8, 255 - mask, 3, cv2.INPAINT_NS)
|
||
painted = mask > 0
|
||
if painted.any() and (~painted).any():
|
||
idx = distance_transform_edt(
|
||
~painted, return_distances=False, return_indices=True,
|
||
)
|
||
inpainted = inpainted[idx[0], idx[1]]
|
||
return inpainted
|