renderer.py

import os
import sys

import numpy as np
import trimesh
from typing import Tuple, Optional

os.environ['PYOPENGL_PLATFORM'] = 'osmesa'

try:
    from OpenGL import osmesa
    if not hasattr(osmesa, 'OSMesaCreateContextAttribs'):
        osmesa.OSMesaCreateContextAttribs = osmesa.OSMesaCreateContext
        print("Patched OSMesaCreateContextAttribs to use OSMesaCreateContext")
except Exception as patch_error:
    print(f"Could not patch OSMesa: {patch_error}")

try:
    import pyrender
except Exception as e:
    print(f"Warning: pyrender import failed with osmesa: {e}")
    try:
        os.environ['PYOPENGL_PLATFORM'] = 'egl'
        import pyrender
        print("Using EGL platform")
    except Exception as e2:
        print(f"Warning: pyrender import failed with egl: {e2}")
        try:
            os.environ['PYOPENGL_PLATFORM'] = 'glx'
            import pyrender
            print("Using GLX platform")
        except Exception as e3:
            print(f"Warning: pyrender import failed with glx: {e3}")
            os.environ['PYOPENGL_PLATFORM'] = 'osmesa'
            import pyrender


class AvatarRenderer:
    def __init__(
        self,
        image_size: int = 512,
        camera_type: str = "orthographic",
        light_intensity: float = 2.0
    ):
        self.image_size = image_size
        self.camera_type = camera_type
        self.light_intensity = light_intensity
        self.renderer = pyrender.OffscreenRenderer(
            viewport_width=image_size,
            viewport_height=image_size,
            point_size=1.0
        )
    
    def render(
        self,
        vertices: np.ndarray,
        faces: np.ndarray,
        camera_pose: Optional[np.ndarray] = None,
        light_pose: Optional[np.ndarray] = None
    ) -> np.ndarray:
        mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
        mesh.vertices -= mesh.vertices.mean(axis=0)
        scale = 1.0 / (mesh.vertices.max(axis=0) - mesh.vertices.min(axis=0)).max()
        mesh.vertices *= scale * 0.8
        
        pyrender_mesh = pyrender.Mesh.from_trimesh(mesh)
        scene = pyrender.Scene(bg_color=[1.0, 1.0, 1.0, 1.0])
        scene.add(pyrender_mesh)
        
        if self.camera_type == "orthographic":
            camera = pyrender.OrthographicCamera(xmag=1.0, ymag=1.0)
        else:
            camera = pyrender.PerspectiveCamera(yfov=np.pi / 3.0, aspectRatio=1.0)
        
        if camera_pose is None:
            camera_pose = np.eye(4)
            camera_pose[:3, 3] = [0, 0, 2.5]
        scene.add(camera, pose=camera_pose)
        
        if light_pose is None:
            light_pose = np.eye(4)
            light_pose[:3, 3] = [0, 0, 2.5]
        
        light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=self.light_intensity)
        scene.add(light, pose=light_pose)
        
        ambient_light = pyrender.SpotLight(
            color=[1.0, 1.0, 1.0],
            intensity=0.5,
            innerConeAngle=np.pi / 16.0,
            outerConeAngle=np.pi / 6.0
        )
        scene.add(ambient_light, pose=light_pose)
        
        color, depth = self.renderer.render(scene)
        
        if color.shape[2] == 4:
            color = color[:, :, :3]
        
        return color
    
    def __del__(self):
        if hasattr(self, 'renderer'):
            self.renderer.delete()


_renderer_instance = None


def get_renderer(image_size: int = 512, camera_type: str = "orthographic") -> AvatarRenderer:
    global _renderer_instance
    if _renderer_instance is None:
        _renderer_instance = AvatarRenderer(image_size=image_size, camera_type=camera_type)
    return _renderer_instance


def render_avatar(
    vertices: np.ndarray,
    faces: np.ndarray,
    image_size: int = 512
) -> np.ndarray:
    renderer = get_renderer(image_size=image_size)
    return renderer.render(vertices, faces)