MeshRenderer.py

import os
os.environ["PYOPENGL_PLATFORM"] = "egl"

from PIL import Image
import numpy as np
import trimesh
import pyrender
from meshlib import mrmeshpy as mm
from meshlib import mrmeshnumpy as mn


class MeshRenderer:
    def __init__(self, num_views=12, resolution=1024):
        """Initialize renderer configuration and lighting/camera defaults.

        Args:
            num_views (int): Number of viewpoints to render around the mesh.
            resolution (int): Square pixel resolution for each rendered view.
        """
        self.num_views = num_views
        self.resolution = resolution
        # self.camera_pose = self.create_camera_pose(radius=2.2, angle=2.0 * np.pi / self.num_views)
        self.key_light_pose = self.create_light_pose(light_pos=[2, 2, 2])
        self.fill_light_pose = self.create_light_pose(light_pos=[-2, 1, 2])
        self.back_light_pose = self.create_light_pose(light_pos=[0, -2, -2])
        self.material = pyrender.MetallicRoughnessMaterial(
        metallicFactor=0.35,
        roughnessFactor=0.5,
        baseColorFactor=[0.85, 0.85, 0.85, 1.0], 
        doubleSided=True
        )
        self.camera = pyrender.PerspectiveCamera(yfov=np.pi / 6.0)
        # self.scene = self.create_scene()
        # self.renderer = pyrender.OffscreenRenderer(self.resolution, self.resolution)

    def create_camera_pose(self, radius, angle, elevation=0.3):
        """Create a 4x4 camera pose matrix for a spherical camera orbit.

        Args:
            radius (float): Distance from the origin to place the camera.
            angle (float): Angle (radians) around the origin.
            elevation (float): Height/elevation of the camera above the XY plane.

        Returns:
            numpy.ndarray: 4x4 homogeneous transform for the camera pose.
        """
        cam_x = radius * np.cos(angle)
        cam_y = radius * np.sin(angle)
        cam_z = elevation

        camera_position = np.array([cam_x, cam_y, cam_z])
        target = np.array([0.0, 0.0, 0.0])
        up = np.array([0.0, 0.0, 1.0])

        forward = target - camera_position
        forward /= np.linalg.norm(forward)

        right = np.cross(forward, up)
        right /= np.linalg.norm(right)

        up = np.cross(right, forward)

        pose = np.eye(4)
        pose[:3, 0] = right
        pose[:3, 1] = up
        pose[:3, 2] = -forward
        pose[:3, 3] = camera_position

        return pose

    def create_light_pose(self, light_pos, target=[0, 0, 0], up=[0, 1, 0]):
        """Build a 4x4 pose matrix that orients a directional light toward a target.

        Args:
            light_pos (sequence): XYZ position of the light source.
            target (sequence): XYZ target the light should point at.
            up (sequence): Up-vector to resolve orientation.

        Returns:
            numpy.ndarray: 4x4 homogeneous transform for the light.
        """
        forward = np.array(target) - np.array(light_pos)
        forward = forward / np.linalg.norm(forward)

        right = np.cross(up, forward)
        right /= np.linalg.norm(right)

        up = np.cross(forward, right)

        rot = np.eye(4)
        rot[:3, 0] = right
        rot[:3, 1] = up
        rot[:3, 2] = -forward

        pose = rot.copy()
        pose[:3, 3] = light_pos

        return pose

    def create_scene(self):
        """Create a new pyrender scene with a 3-point lighting setup.

        Returns:
            pyrender.Scene: Scene populated with directional lights.
        """
        scene = pyrender.Scene(
        bg_color=[30, 30, 30, 255], ambient_light=[0.02, 0.02, 0.02] )

        # 3-point lighting setup
        key_light = pyrender.DirectionalLight(color=np.ones(3), intensity=1.5)
        fill_light = pyrender.DirectionalLight(color=np.ones(3), intensity=0.6)
        back_light = pyrender.DirectionalLight(color=np.ones(3), intensity=0.8)

        scene.add(key_light, pose=self.key_light_pose)
        scene.add(fill_light, pose=self.fill_light_pose)
        scene.add(back_light, pose=self.back_light_pose)
        return scene

    def normalize_mesh(self, mesh):
        """Center and uniformly scale a mesh to fit within a unit cube.

        Args:
            mesh (trimesh.Trimesh or trimesh.Scene): Mesh to normalize.

        Returns:
            trimesh.Trimesh: Normalized single Trimesh instance.
        """
        if isinstance(mesh, trimesh.Scene): # reconcile multi-component meshes into one mesh
            mesh = trimesh.util.concatenate(tuple(mesh.geometry.values())) 
        mesh.apply_translation(-mesh.bounding_box.centroid)
        scale = 1.0 / np.max(mesh.extents)
        mesh.apply_scale(scale)
        return mesh

    def render_views(self, mesh_file, output_folder='renders'):
        """Render multiple views of a mesh and return saved image paths.

        Args:
            mesh_file (str): Path to the mesh file to render.
            output_folder (str): Base output directory to store renders.

        Returns:
            list[str]: Paths to the rendered view images.
        """
        output_folder = os.path.join(output_folder, os.path.splitext(os.path.basename(mesh_file))[0])
        if self.is_mesh_rendered(output_folder):
            print(f"Mesh already rendered. Loading images from {output_folder}...")
            return [os.path.join(output_folder, f) for f in sorted(os.listdir(output_folder))]

        os.environ["PYOPENGL_PLATFORM"] = "egl"
        mesh = trimesh.load(mesh_file, force='mesh')

        if not isinstance(mesh, trimesh.Trimesh):
            raise ValueError(f"Loaded mesh is not a valid Trimesh object: {mesh_file}")
        mesh = self.normalize_mesh(mesh)

        mesh.update_faces(mesh.nondegenerate_faces())
        mesh.update_faces(mesh.unique_faces())
        mesh.remove_unreferenced_vertices()
        mesh.fix_normals()

        render_mesh = pyrender.Mesh.from_trimesh(mesh, material=self.material)

        images = []
        for i in range(self.num_views):
            renderer = pyrender.OffscreenRenderer(self.resolution, self.resolution)
            scene = self.create_scene() # create a new scene for each view to avoid node conflicts
            mesh_node = scene.add(render_mesh)
            angle = i * (2.0 * np.pi / self.num_views)
            camera_pose = self.create_camera_pose(radius=2.2, angle = angle)
            scene.add(self.camera, pose=camera_pose)
            
            color, _ = renderer.render(scene)
            images.append(Image.fromarray(color))

            renderer.delete()

        return self.save_images(images, output_folder)

    def save_images(self, images, output_folder):
        """Save a list of PIL images to disk and return their file paths.

        Args:
            images (list[PIL.Image.Image]): Images to save.
            output_folder (str): Directory to save images into.

        Returns:
            list[str]: Full paths to the saved image files.
        """
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)

        image_paths = []
        for idx, img in enumerate(images):
            filename = os.path.join(output_folder, f"view_{idx:02d}.png")
            img.save(filename) 
            image_paths.append(filename)

        return image_paths

    def is_mesh_rendered(self, output_folder):
        """Return True if a folder contains at least `num_views` files indicating rendered views are present.

        Args:
            output_folder (str): Directory to check for rendered views.

        Returns:
            bool: True when rendered views appear present.
        """
        if not os.path.exists(output_folder):
            return False
        files = os.listdir(output_folder)
        return len(files) >= self.num_views # check if all views were rendered

    def decimate_mesh(self, mesh_file, target_faces=400, output_folder='decimated_meshes'):
        """Decimate a mesh to approximately `target_faces` using meshlib.

        Args:
            mesh_file (str): Path to the input mesh file.
            target_faces (int): Target number of faces after decimation.
            output_folder (str): Directory to save the decimated mesh.

        Returns:
            meshlib mesh: The decimated mesh object.
        """
        mesh = mm.loadMesh(mesh_file)
        mesh.packOptimally()
        settings = mm.DecimateSettings()
        current_faces = mesh.topology.numValidFaces()
        settings.maxDeletedFaces = current_faces - target_faces
        settings.maxError = float('inf')  # no error limit, just reduce to target face count
        settings.subdivideParts = 64 # parallelize decimation across parts for speed

        decimated_mesh = mm.decimateMesh(mesh, settings)
        print(f"Decimated mesh from {current_faces} faces to {mesh.topology.numValidFaces()} faces with error {decimated_mesh.errorIntroduced}.")
        mesh.packOptimally()  # re-pack to clean up after decimation

        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
        output_path = os.path.join(output_folder, os.path.basename(mesh_file))
        mm.saveMesh(mesh, output_path)

        return mesh

    def mesh_to_string(self, mesh_file, decimate=True, target_faces=400):
        """Convert a mesh to a simple text representation (OBJ-like).

        Optionally decimates the mesh first to reduce size.

        Args:
            mesh_file (str): Path to the mesh file.
            decimate (bool): Whether to decimate before conversion.
            target_faces (int): Target face count if decimating.

        Returns:
            str: Multi-line string with 'v' and 'f' lines describing the mesh.
        """
        if decimate:
            mesh = self.decimate_mesh(mesh_file, target_faces=target_faces)
        else:
            mesh = mm.loadMesh(mesh_file)

        vertices = mn.getNumpyVerts(mesh)
        faces = mn.getNumpyFaces(mesh.topology)
        mesh_string = ''

        vertices = self.quantize_vertices(vertices)

        for v in vertices:
            mesh_string += f"v {v[0]} {v[1]} {v[2]}\n"

        for f in faces:
            mesh_string += f"f {f[0]+1} {f[1]+1} {f[2]+1}\n"

        return mesh_string

    def quantize_vertices(self, vertices, bins=64):
        """Quantize vertex coordinates into integer bins.

        Args:
            vertices (numpy.ndarray): Nx3 array of vertex coordinates.
            bins (int): Number of discrete bins per axis.

        Returns:
            numpy.ndarray: Quantized integer vertex coordinates.
        """
        vmin = vertices.min(axis=0)
        vmax = vertices.max(axis=0)

        scale = vmax - vmin
        scale[scale == 0] = 1e-9
        normalized = (vertices - vmin) / scale
        quantized = np.round(normalized * bins).astype(int)
        quantized = np.clip(quantized, 0, bins)

        return quantized