render_room2.py

#!/usr/bin/env python3
"""
Fast 3D room render for WhatsApp virtual background.
Extracts mesh and renders a 2D image without heavy dependencies.
"""

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

# Load the room scene
scene = trimesh.load('room.glb')
print(f"Loaded scene: {type(scene)}")

# Extract the mesh from the scene
if isinstance(scene, trimesh.scene.scene.Scene):
    # Get the first geometry
    geometry_name = list(scene.geometry.keys())[0]
    mesh = scene.geometry[geometry_name]
    print(f"Using geometry: {geometry_name}")
else:
    mesh = scene

print(f"Vertices: {mesh.vertices.shape}")
print(f"Faces: {mesh.faces.shape}")
print(f"Bounds: {mesh.bounds}")

# Compute normals if missing
if mesh.vertex_normals is None or len(mesh.vertex_normals) == 0:
    print("Computing vertex normals...")
    mesh.vertex_normals = mesh.faces_sparse.dot(np.ones(len(mesh.faces))) / 3

# Camera setup (eye-level, looking at the center of the room)
room_center = mesh.centroid
camera_distance = 2.5
eye_height = 1.6

# Place camera at the 'front' of the room (positive Z)
# We look at the center from a slight angle
camera_pos = np.array([0.0, eye_height, camera_distance])
target_pos = room_center
up_vector = np.array([0.0, 1.0, 0.0])

print(f"Camera position: {camera_pos}")
print(f"Target position: {target_pos}")

# Create view matrix using look-at
def normalize(v):
    return v / np.linalg.norm(v)

forward = normalize(target_pos - camera_pos)
right = normalize(np.cross(forward, up_vector))
up = np.cross(right, forward)

# View matrix (row-major for our manual projection)
view_matrix = 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]
])

# Projection parameters
width, height = 1920, 1080
fov_y = math.radians(60)  # Vertical FOV
aspect = width / height
near, far = 0.1, 100.0

# Perspective projection matrix
f = 1.0 / math.tan(fov_y / 2.0)
proj_matrix = np.array([
    [f / aspect, 0, 0, 0],
    [0, f, 0, 0],
    [0, 0, (far + near) / (near - far), (2 * far * near) / (near - far)],
    [0, 0, -1, 0]
])

print("Projection matrix set up")

# Render
img = Image.new('RGB', (width, height), (40, 40, 60))
draw = ImageDraw.Draw(img)

# Process faces: transform, project, and draw
vertices = mesh.vertices
faces = mesh.faces

# Compute face normals for flat shading
face_normals = mesh.face_normals

# We'll pre-compute depth for sorting
face_centers = np.mean(vertices[faces], axis=1)

# Transform centers to view space for depth sorting
centers_view = np.dot(np.hstack([face_centers, np.ones((len(face_centers), 1))]), view_matrix.T)
depths = centers_view[:, 2]

# Sort faces by depth (draw furthest first)
sorted_indices = np.argsort(depths)[::-1]  # Descending depth

print(f"Processing {len(faces)} faces...")

# Light direction (from above and slightly to the right)
light_dir = normalize(np.array([0.5, 1.0, 0.5]))

for i, face_idx in enumerate(sorted_indices[:5000]):  # Limit to 5k faces for speed
    face = faces[face_idx]
    
    # Get vertex positions
    v0 = vertices[face[0]]
    v1 = vertices[face[1]]
    v2 = vertices[face[2]]
    
    # Transform to view space
    def to_view(v):
        p = np.append(v, 1.0)
        p_view = np.dot(view_matrix, p)
        return p_view
    
    v0_view = to_view(v0)
    v1_view = to_view(v1)
    v2_view = to_view(v2)
    
    # Check if all vertices are in front of the near plane
    if v0_view[2] < near and v1_view[2] < near and v2_view[2] < near:
        continue
    
    # Project to clip space
    def project(v):
        p_clip = np.dot(proj_matrix, v)
        if p_clip[3] <= 0:
            return None
        w = p_clip[3]
        x = (p_clip[0] / w) * 0.5 + 0.5
        y = (p_clip[1] / w) * 0.5 + 0.5
        return (int(x * width), int((1 - y) * height))
    
    p0 = project(v0_view)
    p1 = project(v1_view)
    p2 = project(v2_view)
    
    if p0 is None or p1 is None or p2 is None:
        continue
    
    # Simple lighting
    normal = face_normals[face_idx]
    # Transform normal to view space
    normal_view = np.dot(view_matrix[:3, :3], normal)
    light_intensity = max(0.0, np.dot(normalize(normal_view), light_dir))
    
    # Base color with lighting
    base_color = (180, 170, 160)
    shaded_color = tuple(int(c * (0.3 + 0.7 * light_intensity)) for c in base_color)
    
    # Draw triangle (using polygon for filled look, or line for wireframe)
    draw.polygon([p0, p1, p2], fill=shaded_color, outline=None)

# Save the rendered image
img.save('room_render_simple.png')
print("Saved room_render_simple.png")

# Also create a point cloud visualization for debugging
img_pc = Image.new('RGB', (width, height), (10, 10, 20))
draw_pc = ImageDraw.Draw(img_pc)

# Project all vertices
for v in vertices:
    p_view = to_view(v)
    p = project(p_view)
    if p:
        x, y = p
        if 0 <= x < width and 0 <= y < height:
            draw_pc.point((x, y), fill=(200, 200, 200))

img_pc.save('room_pointcloud.png')
print("Saved room_pointcloud.png")