render_room3.py

#!/usr/bin/env python3
"""
Fast 3D room render for WhatsApp virtual background.
Uses software rasterization with numpy/PIL for maximum compatibility.
"""

import trimesh
import numpy as np
from PIL import Image
import math

# Load the room scene
scene = trimesh.load('room.glb')
if isinstance(scene, trimesh.scene.scene.Scene):
    mesh = scene.geometry[list(scene.geometry.keys())[0]]
else:
    mesh = scene

vertices = mesh.vertices.astype(np.float32)
faces = mesh.faces

# Camera setup
room_center = mesh.centroid
camera_pos = np.array([0.0, 1.6, 2.5], dtype=np.float32)
camera_target = room_center
camera_up = np.array([0.0, 1.0, 0.0], dtype=np.float32)

# Build look-at matrix
forward = (camera_target - camera_pos)
forward = forward / np.linalg.norm(forward)
right = np.cross(forward, camera_up)
right = right / np.linalg.norm(right)
up = np.cross(right, forward)

view = np.array([
    [right[0], right[1], right[2], -np.dot(right, camera_pos)],
    [up[0], up[1], up[2], -np.dot(up, camera_pos)],
    [forward[0], forward[1], forward[2], -np.dot(forward, camera_pos)],
    [0, 0, 0, 1]], dtype=np.float32)

# Projection setup
W, H = 1920, 1080
fov = math.radians(60)
f = 1.0 / math.tan(fov / 2)
near, far = 0.1, 100.0

# Transform all vertices
v_homo = np.column_stack([vertices, np.ones(len(vertices), dtype=np.float32)])
v_view = v_homo.dot(view.T)  # Shape (N, 4)

# Project to clip space: x/view_z, y/view_z
vz = v_view[:, 2]
# Keep only vertices in front of the near plane
valid = vz > near
v_proj = np.zeros((len(vertices), 2), dtype=np.float32)
v_proj[valid] = v_view[valid, :2] / v_view[valid, 2:3]

# Screen coordinates
screen_x = (v_proj[:, 0] * f * (W/H) * 0.5 + 0.5) * W
screen_y = (-v_proj[:, 1] * f * 0.5 + 0.5) * H

# Create image buffer
img_arr = np.zeros((H, W, 3), dtype=np.uint8)
img_arr[:, :] = [40, 40, 60]  # Background color

# Simple depth buffer
z_buffer = np.full((H, W), np.inf, dtype=np.float32)

# Compute face normals
face_normals = mesh.face_normals
face_centers = np.mean(vertices[faces], axis=1)
light_dir = np.array([0.5, 1.0, 0.5], dtype=np.float32)
light_dir = light_dir / np.linalg.norm(light_dir)

# Render each face (simple point-based approach for speed)
for i, face in enumerate(faces):
    # Check if face is valid
    v0, v1, v2 = face
    if not (valid[v0] and valid[v1] and valid[v2]):
        continue
    
    # Get screen coordinates
    pts = np.array([
        [screen_x[v0], screen_y[v0]],
        [screen_x[v1], screen_y[v1]],
        [screen_x[v2], screen_y[v2]]
    ], dtype=np.int32)
    
    # Skip if all points are outside
    if (pts[:, 0].min() < 0 or pts[:, 0].max() >= W or
        pts[:, 1].min() < 0 or pts[:, 1].max() >= H):
        continue
    
    # Compute face color based on normal and lighting
    normal = face_normals[i]
    light_intensity = max(0.0, np.dot(normal, light_dir)) * 0.7 + 0.3
    color = (np.array([180, 170, 160], dtype=np.float32) * light_intensity).astype(np.uint8)
    
    # Draw the triangle (using simple scanline for now - just draw 3 points for speed)
    # For a faster approach, we'll just draw vertices and edges
    for pt in pts:
        x, y = pt
        if 0 <= x < W and 0 <= y < H:
            img_arr[y, x] = color
    
    # Draw edges
    for j in range(3):
        x1, y1 = pts[j]
        x2, y2 = pts[(j+1) % 3]
        # Simple line drawing (Bresenham)
        dx = abs(x2 - x1)
        dy = abs(y2 - y1)
        x, y = x1, y1
        sx = 1 if x1 < x2 else -1
        sy = 1 if y1 < y2 else -1
        if dx > dy:
            err = dx // 2
            while x != x2:
                if 0 <= x < W and 0 <= y < H:
                    img_arr[y, x] = color
                err -= dy
                if err < 0:
                    y += sy
                    err += dx
                x += sx
        else:
            err = dy // 2
            while y != y2:
                if 0 <= x < W and 0 <= y < H:
                    img_arr[y, x] = color
                err -= dx
                if err < 0:
                    x += sx
                    err += dy
                y += sy

# Save image
img = Image.fromarray(img_arr)
img.save('room_rendered.png')
print("Saved room_rendered.png")