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
502 lines
18 KiB
Python
502 lines
18 KiB
Python
"""Tests for MLX mesh renderer and camera utilities."""
|
|
|
|
import math
|
|
import sys
|
|
import os
|
|
|
|
import mlx.core as mx
|
|
import numpy as np
|
|
import pytest
|
|
|
|
# Add hy3dpaint to path for imports
|
|
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
|
|
|
|
from DifferentiableRenderer.camera_utils_mlx import (
|
|
get_mv_matrix,
|
|
get_orthographic_projection_matrix,
|
|
get_perspective_projection_matrix,
|
|
transform_pos,
|
|
)
|
|
from DifferentiableRenderer.mesh_render_mlx import (
|
|
MLXRasterizer,
|
|
MeshRenderMLX,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _make_cube_mesh():
|
|
"""A unit cube with outward-pointing normals (CCW winding from outside)."""
|
|
verts = np.array([
|
|
[-0.5, -0.5, -0.5], [0.5, -0.5, -0.5],
|
|
[0.5, 0.5, -0.5], [-0.5, 0.5, -0.5],
|
|
[-0.5, -0.5, 0.5], [0.5, -0.5, 0.5],
|
|
[0.5, 0.5, 0.5], [-0.5, 0.5, 0.5],
|
|
], dtype=np.float32)
|
|
# Winding reversed vs naive ordering so cross(e1,e2) points outward
|
|
faces = np.array([
|
|
[0,2,1],[0,3,2], # front (normal -Z)
|
|
[4,5,6],[4,6,7], # back (normal +Z)
|
|
[0,4,7],[0,7,3], # left (normal -X)
|
|
[1,2,6],[1,6,5], # right (normal +X)
|
|
[3,7,6],[3,6,2], # top (normal +Y)
|
|
[0,1,5],[0,5,4], # bottom (normal -Y)
|
|
], dtype=np.int32)
|
|
return verts, faces
|
|
|
|
|
|
def _make_single_triangle():
|
|
"""A single large triangle in XY plane."""
|
|
verts = np.array([
|
|
[-0.8, -0.8, 0.0],
|
|
[ 0.8, -0.8, 0.0],
|
|
[ 0.0, 0.8, 0.0],
|
|
], dtype=np.float32)
|
|
faces = np.array([[0, 1, 2]], dtype=np.int32)
|
|
return verts, faces
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Camera utilities
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestCameraUtils:
|
|
def test_transform_pos_shape(self):
|
|
mtx = np.eye(4, dtype=np.float32)
|
|
pos = mx.zeros((10, 3))
|
|
result = transform_pos(mtx, pos)
|
|
assert result.shape == (1, 10, 4)
|
|
|
|
def test_transform_pos_keepdim(self):
|
|
mtx = np.eye(4, dtype=np.float32)
|
|
pos = mx.zeros((10, 3))
|
|
result = transform_pos(mtx, pos, keepdim=True)
|
|
assert result.shape == (10, 4)
|
|
|
|
def test_transform_pos_identity(self):
|
|
mtx = np.eye(4, dtype=np.float32)
|
|
pos = mx.array([[1.0, 2.0, 3.0]])
|
|
result = transform_pos(mtx, pos, keepdim=True)
|
|
mx.synchronize()
|
|
expected = [1.0, 2.0, 3.0, 1.0]
|
|
np.testing.assert_allclose(np.array(result[0]), expected, atol=1e-6)
|
|
|
|
def test_transform_pos_with_4d_input(self):
|
|
mtx = np.eye(4, dtype=np.float32)
|
|
pos = mx.array([[1.0, 2.0, 3.0, 1.0]])
|
|
result = transform_pos(mtx, pos, keepdim=True)
|
|
mx.synchronize()
|
|
np.testing.assert_allclose(np.array(result[0]), [1.0, 2.0, 3.0, 1.0], atol=1e-6)
|
|
|
|
def test_mv_matrix_shape(self):
|
|
mv = get_mv_matrix(0, 0, 1.45)
|
|
assert mv.shape == (4, 4)
|
|
assert mv.dtype == np.float32
|
|
|
|
def test_mv_matrix_is_rigid(self):
|
|
mv = get_mv_matrix(30, 45, 2.0)
|
|
R = mv[:3, :3]
|
|
# R^T @ R should be identity for a rotation matrix
|
|
np.testing.assert_allclose(R.T @ R, np.eye(3), atol=1e-5)
|
|
|
|
def test_orthographic_projection(self):
|
|
proj = get_orthographic_projection_matrix()
|
|
assert proj.shape == (4, 4)
|
|
# Origin should map to origin
|
|
result = proj @ np.array([0, 0, -1, 1])
|
|
assert result[3] == 1.0 # w unchanged in ortho
|
|
|
|
def test_perspective_projection(self):
|
|
proj = get_perspective_projection_matrix(49.13, 1.0, 0.01, 100.0)
|
|
assert proj.shape == (4, 4)
|
|
assert proj[3, 2] == -1.0 # perspective divide
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Rasterizer adapter
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestMLXRasterizer:
|
|
def test_rasterize_single_triangle(self):
|
|
verts = mx.array([
|
|
[-0.5, -0.5, 0.5, 1.0],
|
|
[ 0.5, -0.5, 0.5, 1.0],
|
|
[ 0.0, 0.5, 0.5, 1.0],
|
|
], dtype=mx.float32)
|
|
faces = mx.array([[0, 1, 2]], dtype=mx.int32)
|
|
|
|
fi, bary = MLXRasterizer.rasterize(
|
|
mx.expand_dims(verts, 0), faces, (16, 16)
|
|
)
|
|
mx.synchronize()
|
|
assert fi.shape == (16, 16)
|
|
assert bary.shape == (16, 16, 3)
|
|
covered = (fi > 0).astype(mx.int32).sum().item()
|
|
assert covered > 0
|
|
|
|
def test_interpolate_shape(self):
|
|
verts = mx.array([
|
|
[-0.5, -0.5, 0.5, 1.0],
|
|
[ 0.5, -0.5, 0.5, 1.0],
|
|
[ 0.0, 0.5, 0.5, 1.0],
|
|
], dtype=mx.float32)
|
|
faces = mx.array([[0, 1, 2]], dtype=mx.int32)
|
|
|
|
fi, bary = MLXRasterizer.rasterize(
|
|
mx.expand_dims(verts, 0), faces, (16, 16)
|
|
)
|
|
colors = mx.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=mx.float32)
|
|
result = MLXRasterizer.interpolate(
|
|
mx.expand_dims(colors, 0), fi, bary, faces
|
|
)
|
|
mx.synchronize()
|
|
assert result.shape == (1, 16, 16, 3)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# MeshRenderMLX
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestMeshRenderMLX:
|
|
def setup_method(self):
|
|
self.renderer = MeshRenderMLX(
|
|
default_resolution=32,
|
|
shader_type="face",
|
|
)
|
|
verts, faces = _make_cube_mesh()
|
|
self.renderer.set_mesh(verts, faces)
|
|
|
|
def test_render_normal_shape(self):
|
|
result = self.renderer.render_normal(0, 0)
|
|
assert result.shape == (32, 32, 3)
|
|
assert result.dtype == np.float32
|
|
|
|
def test_render_normal_coverage(self):
|
|
result = self.renderer.render_normal(0, 0, bg_color=(1, 1, 1))
|
|
# The cube should cover a portion of the image
|
|
# Background is (1,1,1) after normalization
|
|
# Non-background pixels should differ from pure white
|
|
bg_pixel = np.array([1.0, 1.0, 1.0])
|
|
not_bg = np.any(np.abs(result - bg_pixel) > 0.01, axis=-1)
|
|
coverage = not_bg.sum()
|
|
total = 32 * 32
|
|
assert coverage > total * 0.05, f"Too little normal coverage: {coverage}/{total}"
|
|
|
|
def test_render_position_shape(self):
|
|
result = self.renderer.render_position(0, 0)
|
|
assert result.shape == (32, 32, 3)
|
|
|
|
def test_render_position_values_in_range(self):
|
|
result = self.renderer.render_position(0, 0, bg_color=(1, 1, 1))
|
|
# Position values should be in [0, 1] range
|
|
assert result.min() >= -0.01
|
|
assert result.max() <= 1.01
|
|
|
|
def test_render_alpha_shape(self):
|
|
result = self.renderer.render_alpha(0, 0)
|
|
# Returns (1, H, W, 1) face indices matching original MeshRender
|
|
assert result.shape == (1, 32, 32, 1)
|
|
|
|
def test_render_alpha_face_indices(self):
|
|
result = self.renderer.render_alpha(0, 0)
|
|
# Should contain face indices: 0 = background, >0 = face IDs
|
|
assert result.dtype == np.int64
|
|
assert (result >= 0).all()
|
|
|
|
def test_render_alpha_has_coverage(self):
|
|
result = self.renderer.render_alpha(0, 0)
|
|
covered = (result > 0).sum()
|
|
assert covered > 0, "Cube should be visible from front"
|
|
|
|
def test_render_from_different_angles(self):
|
|
"""Verify rendering from different elevations produces different results."""
|
|
# Use elevation change (not azimuth) since a cube is rotationally
|
|
# symmetric under 90-degree azimuth steps with face shading.
|
|
n1 = self.renderer.render_normal(0, 0)
|
|
n2 = self.renderer.render_normal(45, 0)
|
|
diff = np.abs(n1 - n2).sum()
|
|
assert diff > 1.0, "Different elevations should produce different normals"
|
|
|
|
|
|
class TestMeshRenderMLXTriangle:
|
|
def setup_method(self):
|
|
self.renderer = MeshRenderMLX(
|
|
default_resolution=16,
|
|
shader_type="face",
|
|
)
|
|
verts, faces = _make_single_triangle()
|
|
self.renderer.set_mesh(verts, faces, auto_center=True)
|
|
|
|
def test_render_normal_nonzero(self):
|
|
result = self.renderer.render_normal(0, 0)
|
|
assert result.shape == (16, 16, 3)
|
|
# Should have some non-background pixels
|
|
bg = np.array([1.0, 1.0, 1.0])
|
|
not_bg = np.any(np.abs(result - bg) > 0.01, axis=-1)
|
|
assert not_bg.sum() > 0
|
|
|
|
|
|
class TestMeshNormalization:
|
|
def test_auto_center(self):
|
|
renderer = MeshRenderMLX(default_resolution=16)
|
|
verts = np.array([
|
|
[10.0, 10.0, 10.0],
|
|
[12.0, 10.0, 10.0],
|
|
[11.0, 12.0, 10.0],
|
|
], dtype=np.float32)
|
|
faces = np.array([[0, 1, 2]], dtype=np.int32)
|
|
renderer.set_mesh(verts, faces, auto_center=True)
|
|
|
|
# After normalization, mesh should be centered near origin
|
|
vtx = np.array(renderer.vtx_pos)
|
|
center = vtx.mean(axis=0)
|
|
np.testing.assert_allclose(center, [0, 0, 0], atol=0.5)
|
|
|
|
def test_no_auto_center(self):
|
|
renderer = MeshRenderMLX(default_resolution=16)
|
|
verts, faces = _make_cube_mesh()
|
|
renderer.set_mesh(verts, faces, auto_center=False)
|
|
assert renderer.scale_factor == 1.0
|
|
|
|
|
|
def _make_cube_with_uvs():
|
|
"""A cube with simple UV coordinates."""
|
|
verts, faces = _make_cube_mesh()
|
|
# Simple planar UV mapping (use xy projected to [0,1])
|
|
uv = np.zeros((len(verts), 2), dtype=np.float32)
|
|
uv[:, 0] = (verts[:, 0] - verts[:, 0].min()) / (verts[:, 0].max() - verts[:, 0].min())
|
|
uv[:, 1] = (verts[:, 1] - verts[:, 1].min()) / (verts[:, 1].max() - verts[:, 1].min())
|
|
return verts, faces, uv
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Texture baking
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestExtractTextiles:
|
|
def test_textiles_populated(self):
|
|
renderer = MeshRenderMLX(default_resolution=32, texture_size=32)
|
|
verts, faces, uv = _make_cube_with_uvs()
|
|
renderer.set_mesh(verts, faces, vtx_uv=uv, uv_idx=faces)
|
|
assert renderer.tex_position is not None
|
|
assert renderer.tex_grid is not None
|
|
assert len(renderer.tex_position) > 0
|
|
|
|
def test_textiles_shape(self):
|
|
renderer = MeshRenderMLX(default_resolution=32, texture_size=16)
|
|
verts, faces, uv = _make_cube_with_uvs()
|
|
renderer.set_mesh(verts, faces, vtx_uv=uv, uv_idx=faces)
|
|
k = len(renderer.tex_position)
|
|
assert renderer.tex_position.shape == (k, 4)
|
|
assert renderer.tex_normal.shape == (k, 3)
|
|
assert renderer.tex_grid.shape == (k, 2)
|
|
assert renderer.texture_indices.shape == (16, 16)
|
|
|
|
|
|
class TestUvFeatureMap:
|
|
def test_shape(self):
|
|
renderer = MeshRenderMLX(default_resolution=32, texture_size=16)
|
|
verts, faces, uv = _make_cube_with_uvs()
|
|
renderer.set_mesh(verts, faces, vtx_uv=uv, uv_idx=faces)
|
|
feat = mx.ones((len(verts), 3))
|
|
result = renderer.uv_feature_map(feat)
|
|
mx.synchronize()
|
|
assert result.shape == (16, 16, 3)
|
|
|
|
|
|
class TestBackProject:
|
|
def setup_method(self):
|
|
self.renderer = MeshRenderMLX(
|
|
default_resolution=32, texture_size=32, boundary_scale=0,
|
|
)
|
|
verts, faces, uv = _make_cube_with_uvs()
|
|
self.renderer.set_mesh(verts, faces, vtx_uv=uv, uv_idx=faces)
|
|
|
|
def test_output_shapes(self):
|
|
# Render a solid red image and back-project
|
|
image = np.ones((32, 32, 3), dtype=np.float32) * np.array([1, 0, 0])
|
|
texture, cos_map, boundary = self.renderer.back_project(image, 0, 0)
|
|
assert texture.shape == (32, 32, 3)
|
|
assert cos_map.shape == (32, 32, 1)
|
|
assert boundary.shape == (32, 32, 1)
|
|
|
|
def test_nonzero_texture(self):
|
|
image = np.ones((32, 32, 3), dtype=np.float32) * 0.5
|
|
texture, cos_map, _ = self.renderer.back_project(image, 0, 0)
|
|
# Some texture pixels should be filled
|
|
filled = (cos_map > 0).sum()
|
|
assert filled > 0, "back_project should fill some texture pixels"
|
|
|
|
|
|
class TestBakeTexture:
|
|
def test_multi_view_bake(self):
|
|
renderer = MeshRenderMLX(
|
|
default_resolution=32, texture_size=32, boundary_scale=0,
|
|
)
|
|
verts, faces, uv = _make_cube_with_uvs()
|
|
renderer.set_mesh(verts, faces, vtx_uv=uv, uv_idx=faces)
|
|
|
|
# Bake from 2 views with solid colors
|
|
img1 = np.ones((32, 32, 3), dtype=np.float32) * 0.8
|
|
img2 = np.ones((32, 32, 3), dtype=np.float32) * 0.6
|
|
|
|
texture, trust = renderer.bake_texture(
|
|
[img1, img2], [0, 0], [0, 90], exp=1,
|
|
)
|
|
assert texture.shape == (32, 32, 3)
|
|
assert trust.shape == (32, 32, 1)
|
|
filled = trust.sum()
|
|
assert filled > 0, "Baked texture should have some coverage"
|
|
|
|
def test_coverage_increases_with_views(self):
|
|
renderer = MeshRenderMLX(
|
|
default_resolution=32, texture_size=32, boundary_scale=0,
|
|
)
|
|
verts, faces, uv = _make_cube_with_uvs()
|
|
renderer.set_mesh(verts, faces, vtx_uv=uv, uv_idx=faces)
|
|
|
|
img = np.ones((32, 32, 3), dtype=np.float32) * 0.5
|
|
|
|
_, trust1 = renderer.bake_texture([img], [0], [0], exp=1)
|
|
_, trust2 = renderer.bake_texture(
|
|
[img, img], [0, 0], [0, 90], exp=1,
|
|
)
|
|
c1 = trust1.sum()
|
|
c2 = trust2.sum()
|
|
assert c2 >= c1, "More views should give equal or more coverage"
|
|
|
|
|
|
class TestFastBakeTexture:
|
|
def test_weighted_merge(self):
|
|
renderer = MeshRenderMLX(texture_size=4)
|
|
# Two views covering different halves: no 99% overlap skip
|
|
t1 = np.ones((4, 4, 3), dtype=np.float32)
|
|
t2 = np.ones((4, 4, 3), dtype=np.float32) * 0.5
|
|
c1 = np.zeros((4, 4, 1), dtype=np.float32)
|
|
c2 = np.zeros((4, 4, 1), dtype=np.float32)
|
|
c1[:2, :] = 1.0 # top half
|
|
c2[2:, :] = 1.0 # bottom half
|
|
merged, trust = renderer.fast_bake_texture([t1, t2], [c1, c2])
|
|
# Top half = 1.0, bottom half = 0.5
|
|
np.testing.assert_allclose(merged[:2], 1.0, atol=1e-6)
|
|
np.testing.assert_allclose(merged[2:], 0.5, atol=1e-6)
|
|
assert trust.all()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Pipeline integration (ViewProcessorMLX + MeshRenderMLX)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class TestMeshManagement:
|
|
def test_get_face_areas(self):
|
|
renderer = MeshRenderMLX(default_resolution=16)
|
|
verts, faces = _make_cube_mesh()
|
|
renderer.set_mesh(verts, faces)
|
|
areas = renderer.get_face_areas()
|
|
assert areas.shape == (12,)
|
|
assert (areas > 0).all()
|
|
|
|
def test_get_face_areas_one_indexed(self):
|
|
renderer = MeshRenderMLX(default_resolution=16)
|
|
verts, faces = _make_cube_mesh()
|
|
renderer.set_mesh(verts, faces)
|
|
areas = renderer.get_face_areas(from_one_index=True)
|
|
assert areas.shape == (13,) # 12 faces + 1 padding
|
|
assert areas[0] == 0.0
|
|
|
|
def test_render_alpha_returns_face_indices(self):
|
|
renderer = MeshRenderMLX(default_resolution=32)
|
|
verts, faces = _make_cube_mesh()
|
|
renderer.set_mesh(verts, faces)
|
|
alpha = renderer.render_alpha(0, 0)
|
|
assert alpha.shape == (1, 32, 32, 1)
|
|
assert alpha.dtype == np.int64
|
|
# Should have face indices > 0 for visible faces
|
|
unique = np.unique(alpha)
|
|
assert 0 in unique # background
|
|
assert len(unique) > 1 # at least one face visible
|
|
|
|
def test_set_get_texture(self):
|
|
renderer = MeshRenderMLX(default_resolution=16, texture_size=8)
|
|
tex = np.random.rand(8, 8, 3).astype(np.float32)
|
|
renderer.set_texture(tex)
|
|
got = renderer.get_texture()
|
|
np.testing.assert_array_equal(got, tex)
|
|
|
|
def test_get_mesh_roundtrip(self):
|
|
renderer = MeshRenderMLX(default_resolution=16)
|
|
verts, faces = _make_cube_mesh()
|
|
renderer.set_mesh(verts, faces, auto_center=False)
|
|
vtx_out, idx_out, _, _ = renderer.get_mesh(normalize=True)
|
|
# After set_mesh + get_mesh, geometry should survive the roundtrip
|
|
assert vtx_out.shape == verts.shape
|
|
assert idx_out.shape == faces.shape
|
|
|
|
def test_render_normal_pil(self):
|
|
from PIL import Image as PILImage
|
|
renderer = MeshRenderMLX(default_resolution=16)
|
|
verts, faces = _make_cube_mesh()
|
|
renderer.set_mesh(verts, faces)
|
|
result = renderer.render_normal(0, 0, return_type="pl")
|
|
assert isinstance(result, PILImage.Image)
|
|
assert result.size == (16, 16)
|
|
|
|
def test_set_default_render_resolution(self):
|
|
renderer = MeshRenderMLX(default_resolution=32)
|
|
assert renderer.default_resolution == (32, 32)
|
|
renderer.set_default_render_resolution(64)
|
|
assert renderer.default_resolution == (64, 64)
|
|
|
|
|
|
class TestViewProcessorMLX:
|
|
def setup_method(self):
|
|
from utils.pipeline_utils_mlx import ViewProcessorMLX
|
|
|
|
class FakeConfig:
|
|
bake_exp = 4
|
|
|
|
self.renderer = MeshRenderMLX(
|
|
default_resolution=32, texture_size=32, boundary_scale=0,
|
|
)
|
|
verts, faces, uv = _make_cube_with_uvs()
|
|
self.renderer.set_mesh(verts, faces, vtx_uv=uv, uv_idx=faces)
|
|
self.vp = ViewProcessorMLX(FakeConfig(), self.renderer)
|
|
|
|
def test_render_normal_multiview(self):
|
|
from PIL import Image as PILImage
|
|
normals = self.vp.render_normal_multiview([0, 0], [0, 90])
|
|
assert len(normals) == 2
|
|
assert isinstance(normals[0], PILImage.Image)
|
|
|
|
def test_render_position_multiview(self):
|
|
from PIL import Image as PILImage
|
|
positions = self.vp.render_position_multiview([0], [0])
|
|
assert len(positions) == 1
|
|
assert isinstance(positions[0], PILImage.Image)
|
|
|
|
def test_bake_from_multiview(self):
|
|
img1 = np.ones((32, 32, 3), dtype=np.float32) * 0.7
|
|
img2 = np.ones((32, 32, 3), dtype=np.float32) * 0.3
|
|
texture, trust = self.vp.bake_from_multiview(
|
|
[img1, img2], [0, 0], [0, 90], [1.0, 1.0],
|
|
)
|
|
assert texture.shape == (32, 32, 3)
|
|
assert trust.shape == (32, 32, 1)
|
|
|
|
def test_bake_view_selection(self):
|
|
elevs = [0, 0, 0, 0, 90, -90]
|
|
azims = [0, 90, 180, 270, 0, 180]
|
|
weights = [1, 0.1, 0.5, 0.1, 0.05, 0.05]
|
|
sel_e, sel_a, sel_w = self.vp.bake_view_selection(
|
|
elevs, azims, weights, max_selected_view_num=6,
|
|
)
|
|
assert len(sel_e) == 6
|
|
assert len(sel_a) == 6
|
|
assert len(sel_w) == 6
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|