149 lines
5.3 KiB
Python
149 lines
5.3 KiB
Python
"""
|
|
Rasterization parity test for Trellis2 texture baking.
|
|
Tests the exact UV-space rasterization scenario used by postprocess.py:
|
|
- Positions are UV coords mapped to clip space (z=0, w=1)
|
|
- All triangles share the same depth
|
|
- Each face must get its correct unique ID
|
|
|
|
This is a focused unit test for the Metal rasterizer in the Trellis2 context.
|
|
For full E2E parity (Metal vs CPU postprocess), see test_postprocess_parity.py.
|
|
"""
|
|
import torch
|
|
import numpy as np
|
|
import struct
|
|
import pytest
|
|
|
|
|
|
def float_to_triidx(f):
|
|
if f <= 16777216.0:
|
|
return int(f)
|
|
bits = struct.unpack('I', struct.pack('f', f))[0]
|
|
return bits - 0x4a800000
|
|
|
|
|
|
@pytest.fixture
|
|
def mtl_rast():
|
|
try:
|
|
from mtldiffrast.torch.ops import MtlRasterizeContext, rasterize
|
|
ctx = MtlRasterizeContext()
|
|
return ctx, rasterize
|
|
except ImportError:
|
|
pytest.skip("mtldiffrast not built")
|
|
|
|
|
|
class TestTrellis2RastParity:
|
|
"""UV-space rasterization as used by Trellis2 postprocess.py."""
|
|
|
|
def test_uv_quad_mesh(self, mtl_rast):
|
|
"""Simple UV quad — the minimal texture baking scenario."""
|
|
ctx, rasterize = mtl_rast
|
|
|
|
# UV coords mapped to clip space: uv * 2 - 1, z=0, w=1
|
|
pos = torch.tensor([
|
|
[-1.0, -1.0, 0.0, 1.0],
|
|
[ 1.0, -1.0, 0.0, 1.0],
|
|
[ 1.0, 1.0, 0.0, 1.0],
|
|
[-1.0, 1.0, 0.0, 1.0],
|
|
], dtype=torch.float32)
|
|
tri = torch.tensor([[0, 1, 2], [0, 2, 3]], dtype=torch.int32)
|
|
|
|
rast_out, _ = rasterize(ctx, pos, tri, resolution=[512, 512])
|
|
rast_np = rast_out[0].numpy()
|
|
|
|
# Full coverage
|
|
covered = (rast_np[:, :, 3] != 0).sum()
|
|
assert covered == 512 * 512, f"Expected full coverage, got {covered}"
|
|
|
|
# Both face IDs present
|
|
face_ids = set()
|
|
for py in range(0, 512, 32):
|
|
for px in range(0, 512, 32):
|
|
if rast_np[py, px, 3] != 0:
|
|
fid = float_to_triidx(rast_np[py, px, 3]) - 1
|
|
face_ids.add(fid)
|
|
assert face_ids == {0, 1}, f"Expected face IDs {{0, 1}}, got {face_ids}"
|
|
|
|
def test_random_uv_mesh(self, mtl_rast):
|
|
"""Random UV layout with many triangles — simulates real Trellis2 mesh."""
|
|
ctx, rasterize = mtl_rast
|
|
torch.manual_seed(42)
|
|
|
|
# Generate a grid mesh with random UV perturbation
|
|
N = 20 # 20x20 grid = 800 triangles
|
|
verts = []
|
|
for y in range(N + 1):
|
|
for x in range(N + 1):
|
|
u = x / N
|
|
v = y / N
|
|
# Small random perturbation to avoid perfectly uniform spacing
|
|
if 0 < x < N and 0 < y < N:
|
|
u += (torch.rand(1).item() - 0.5) * 0.02
|
|
v += (torch.rand(1).item() - 0.5) * 0.02
|
|
cx = u * 2.0 - 1.0
|
|
cy = v * 2.0 - 1.0
|
|
verts.append([cx, cy, 0.0, 1.0])
|
|
|
|
tris = []
|
|
for y in range(N):
|
|
for x in range(N):
|
|
i = y * (N + 1) + x
|
|
tris.append([i, i + 1, i + N + 2])
|
|
tris.append([i, i + N + 2, i + N + 1])
|
|
|
|
pos = torch.tensor(verts, dtype=torch.float32)
|
|
tri = torch.tensor(tris, dtype=torch.int32)
|
|
|
|
rast_out, _ = rasterize(ctx, pos, tri, resolution=[256, 256])
|
|
rast_np = rast_out[0].numpy()
|
|
|
|
# Should have high coverage
|
|
covered = (rast_np[:, :, 3] != 0).sum()
|
|
total = 256 * 256
|
|
coverage_pct = covered / total * 100
|
|
assert coverage_pct > 90, f"Expected >90% coverage, got {coverage_pct:.1f}%"
|
|
|
|
# Should have many distinct face IDs (800 total)
|
|
face_ids = set()
|
|
for py in range(256):
|
|
for px in range(256):
|
|
if rast_np[py, px, 3] != 0:
|
|
fid = float_to_triidx(rast_np[py, px, 3]) - 1
|
|
face_ids.add(fid)
|
|
# At 256x256 resolution, most of the 800 triangles should appear
|
|
assert len(face_ids) > 600, f"Expected >600 distinct face IDs, got {len(face_ids)}"
|
|
|
|
def test_barycentrics_interpolation_valid(self, mtl_rast):
|
|
"""Verify interpolated barycentrics are valid for texture sampling."""
|
|
ctx, rasterize = mtl_rast
|
|
from mtldiffrast.torch.ops import interpolate
|
|
|
|
pos = torch.tensor([
|
|
[-1.0, -1.0, 0.0, 1.0],
|
|
[ 1.0, -1.0, 0.0, 1.0],
|
|
[ 1.0, 1.0, 0.0, 1.0],
|
|
[-1.0, 1.0, 0.0, 1.0],
|
|
], dtype=torch.float32)
|
|
tri = torch.tensor([[0, 1, 2], [0, 2, 3]], dtype=torch.int32)
|
|
# UV coordinates as vertex attributes
|
|
uv_attr = torch.tensor([
|
|
[0.0, 0.0, 0.0],
|
|
[1.0, 0.0, 0.0],
|
|
[1.0, 1.0, 0.0],
|
|
[0.0, 1.0, 0.0],
|
|
], dtype=torch.float32)
|
|
|
|
rast_out, _ = rasterize(ctx, pos, tri, resolution=[64, 64])
|
|
interp_out, _ = interpolate(uv_attr, rast_out, tri)
|
|
interp_np = interp_out[0].numpy()
|
|
|
|
# Interpolated UVs should be in [0, 1]
|
|
mask = rast_out[0, :, :, 3].numpy() != 0
|
|
u_vals = interp_np[mask, 0]
|
|
v_vals = interp_np[mask, 1]
|
|
assert (u_vals >= -0.01).all() and (u_vals <= 1.01).all(), f"u out of range: [{u_vals.min()}, {u_vals.max()}]"
|
|
assert (v_vals >= -0.01).all() and (v_vals <= 1.01).all(), f"v out of range: [{v_vals.min()}, {v_vals.max()}]"
|
|
|
|
|
|
if __name__ == '__main__':
|
|
pytest.main([__file__, '-v'])
|