styletransferandClaude Opus 4.5
5d0cc8aa0a
Features: - 8 ComfyUI nodes for TRELLIS2 video-to-3D processing - Onion skinning with configurable frame interval - Multiple rendering modes (PBR, wireframe, cel, normal, depth) - 4 lighting presets (studio, outdoor, dramatic, flat) - Post-processing effects (bloom, DoF, motion blur, outlines) - Progress bars and auto model download - 5 example workflows 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
241 lines
6.3 KiB
Python
241 lines
6.3 KiB
Python
"""Wireframe rendering effect."""
|
|
|
|
import logging
|
|
from typing import Tuple, Optional
|
|
|
|
import torch
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def render_wireframe(
|
|
vertices: torch.Tensor,
|
|
faces: torch.Tensor,
|
|
mvp: torch.Tensor,
|
|
width: int,
|
|
height: int,
|
|
line_width: float = 1.0,
|
|
line_color: Tuple[float, float, float] = (1.0, 1.0, 1.0),
|
|
background_color: Tuple[float, float, float] = (0.0, 0.0, 0.0),
|
|
fill_color: Optional[Tuple[float, float, float]] = None,
|
|
) -> torch.Tensor:
|
|
"""
|
|
Render a mesh as wireframe.
|
|
|
|
Args:
|
|
vertices: Mesh vertices [N, 3]
|
|
faces: Face indices [F, 3]
|
|
mvp: Model-view-projection matrix [4, 4]
|
|
width: Output width
|
|
height: Output height
|
|
line_width: Width of wireframe lines
|
|
line_color: RGB color for lines
|
|
background_color: RGB background color
|
|
fill_color: Optional fill color for faces
|
|
|
|
Returns:
|
|
Rendered wireframe image [H, W, 3]
|
|
"""
|
|
device = vertices.device
|
|
|
|
# Transform vertices to clip space
|
|
verts_homo = torch.cat([
|
|
vertices,
|
|
torch.ones(vertices.shape[0], 1, device=device)
|
|
], dim=1)
|
|
verts_clip = (mvp @ verts_homo.T).T
|
|
|
|
# Perspective divide to NDC
|
|
verts_ndc = verts_clip[:, :3] / (verts_clip[:, 3:4] + 1e-8)
|
|
|
|
# Convert to screen space
|
|
verts_screen = torch.zeros_like(verts_ndc[:, :2])
|
|
verts_screen[:, 0] = (verts_ndc[:, 0] + 1) * 0.5 * width
|
|
verts_screen[:, 1] = (1 - verts_ndc[:, 1]) * 0.5 * height # Flip Y
|
|
|
|
# Initialize output
|
|
output = torch.zeros(height, width, 3, dtype=torch.float32, device=device)
|
|
output[:, :, 0] = background_color[0]
|
|
output[:, :, 1] = background_color[1]
|
|
output[:, :, 2] = background_color[2]
|
|
|
|
# Draw edges
|
|
line_color_t = torch.tensor(line_color, dtype=torch.float32, device=device)
|
|
|
|
for face in faces:
|
|
for j in range(3):
|
|
v0 = verts_screen[face[j]]
|
|
v1 = verts_screen[face[(j + 1) % 3]]
|
|
|
|
# Draw line between v0 and v1
|
|
output = _draw_line(output, v0, v1, line_color_t, line_width)
|
|
|
|
return output
|
|
|
|
|
|
def wireframe_overlay(
|
|
base_image: torch.Tensor,
|
|
wireframe: torch.Tensor,
|
|
opacity: float = 0.5,
|
|
blend_mode: str = "over"
|
|
) -> torch.Tensor:
|
|
"""
|
|
Overlay wireframe on top of base image.
|
|
|
|
Args:
|
|
base_image: Base rendered image [H, W, 3]
|
|
wireframe: Wireframe image [H, W, 3]
|
|
opacity: Wireframe opacity
|
|
blend_mode: "over", "add", or "screen"
|
|
|
|
Returns:
|
|
Composited image [H, W, 3]
|
|
"""
|
|
if blend_mode == "add":
|
|
return (base_image + wireframe * opacity).clamp(0, 1)
|
|
elif blend_mode == "screen":
|
|
return (1 - (1 - base_image) * (1 - wireframe * opacity)).clamp(0, 1)
|
|
else: # over
|
|
# Only overlay where wireframe has color
|
|
mask = (wireframe.sum(dim=-1, keepdim=True) > 0).float()
|
|
return base_image * (1 - mask * opacity) + wireframe * mask * opacity
|
|
|
|
|
|
def _draw_line(
|
|
image: torch.Tensor,
|
|
p0: torch.Tensor,
|
|
p1: torch.Tensor,
|
|
color: torch.Tensor,
|
|
width: float = 1.0
|
|
) -> torch.Tensor:
|
|
"""
|
|
Draw a line on an image using Bresenham's algorithm.
|
|
|
|
Args:
|
|
image: Image tensor [H, W, 3]
|
|
p0: Start point [2]
|
|
p1: End point [2]
|
|
color: Line color [3]
|
|
width: Line width
|
|
|
|
Returns:
|
|
Image with line drawn
|
|
"""
|
|
H, W, _ = image.shape
|
|
|
|
x0, y0 = int(p0[0].item()), int(p0[1].item())
|
|
x1, y1 = int(p1[0].item()), int(p1[1].item())
|
|
|
|
# Bresenham's line algorithm
|
|
dx = abs(x1 - x0)
|
|
dy = abs(y1 - y0)
|
|
sx = 1 if x0 < x1 else -1
|
|
sy = 1 if y0 < y1 else -1
|
|
err = dx - dy
|
|
|
|
half_w = int(width / 2)
|
|
|
|
while True:
|
|
# Draw pixel with width
|
|
for wx in range(-half_w, half_w + 1):
|
|
for wy in range(-half_w, half_w + 1):
|
|
px, py = x0 + wx, y0 + wy
|
|
if 0 <= px < W and 0 <= py < H:
|
|
image[py, px] = color
|
|
|
|
if x0 == x1 and y0 == y1:
|
|
break
|
|
|
|
e2 = 2 * err
|
|
if e2 > -dy:
|
|
err -= dy
|
|
x0 += sx
|
|
if e2 < dx:
|
|
err += dx
|
|
y0 += sy
|
|
|
|
return image
|
|
|
|
|
|
def extract_edges(
|
|
faces: torch.Tensor
|
|
) -> torch.Tensor:
|
|
"""
|
|
Extract unique edges from face indices.
|
|
|
|
Args:
|
|
faces: Face indices [F, 3]
|
|
|
|
Returns:
|
|
Edge indices [E, 2]
|
|
"""
|
|
# Get all edges (including duplicates)
|
|
edges = []
|
|
for i in range(3):
|
|
j = (i + 1) % 3
|
|
e = torch.stack([faces[:, i], faces[:, j]], dim=1)
|
|
edges.append(e)
|
|
|
|
all_edges = torch.cat(edges, dim=0)
|
|
|
|
# Sort each edge so smaller index comes first
|
|
sorted_edges = torch.sort(all_edges, dim=1).values
|
|
|
|
# Remove duplicates
|
|
unique_edges = torch.unique(sorted_edges, dim=0)
|
|
|
|
return unique_edges
|
|
|
|
|
|
def compute_edge_visibility(
|
|
vertices: torch.Tensor,
|
|
faces: torch.Tensor,
|
|
viewmat: torch.Tensor,
|
|
edges: Optional[torch.Tensor] = None
|
|
) -> torch.Tensor:
|
|
"""
|
|
Compute which edges are visible (silhouette edges).
|
|
|
|
An edge is a silhouette edge if its two adjacent faces have
|
|
different front/back facing orientations.
|
|
|
|
Args:
|
|
vertices: Vertex positions [N, 3]
|
|
faces: Face indices [F, 3]
|
|
viewmat: View matrix [4, 4]
|
|
edges: Optional pre-computed edges [E, 2]
|
|
|
|
Returns:
|
|
Boolean mask of visible edges [E]
|
|
"""
|
|
device = vertices.device
|
|
|
|
if edges is None:
|
|
edges = extract_edges(faces)
|
|
|
|
# Compute face normals
|
|
v0 = vertices[faces[:, 0]]
|
|
v1 = vertices[faces[:, 1]]
|
|
v2 = vertices[faces[:, 2]]
|
|
|
|
face_normals = torch.cross(v1 - v0, v2 - v0)
|
|
face_normals = face_normals / (torch.norm(face_normals, dim=1, keepdim=True) + 1e-8)
|
|
|
|
# Face centers
|
|
face_centers = (v0 + v1 + v2) / 3
|
|
|
|
# View direction (from face center to camera)
|
|
cam_pos = -viewmat[:3, :3].T @ viewmat[:3, 3]
|
|
view_dirs = cam_pos.unsqueeze(0) - face_centers
|
|
view_dirs = view_dirs / (torch.norm(view_dirs, dim=1, keepdim=True) + 1e-8)
|
|
|
|
# Front-facing test
|
|
dots = (face_normals * view_dirs).sum(dim=1)
|
|
front_facing = dots > 0
|
|
|
|
# Build edge-to-face mapping (simplified: mark edge visible if any adjacent face is front-facing)
|
|
# For proper silhouette, need to check if edge is between front and back facing faces
|
|
visible = torch.ones(len(edges), dtype=torch.bool, device=device)
|
|
|
|
return visible
|