step-parse/skill.src/modules/renderer.py

"""
renderer.py — Offscreen PNG thumbnail generation.

Pipeline (with color): build123d → GLTF export → trimesh scene → pyrender → PNG
Pipeline (fallback):   build123d → per-solid STL → colored trimesh scene → pyrender → PNG

Coordinate convention AFTER STEP→GLTF export→trimesh load:
  trimesh applies a Z-up→Y-up GLTF convention that swaps STEP's Y and Z axes:
    X = width  (~248mm for MR16) — left/right
    Y = depth  (~41mm  for MR16) — front/back; Y_min=screen face, Y_max=back panel
    Z = height (~459mm for MR16) — tall axis;  Z_min=top end,    Z_max=bottom end
  → "front" camera sits at -Y (screen side) looking toward +Y to see the LCD face.
  → World "up" vector is (0,0,-1) — negative Z = top of display in image.

6 standard views: front, rear, left, right, iso_left, iso_right
"""
import logging
import tempfile
from pathlib import Path

import numpy as np

from .loader import StepModel

logger = logging.getLogger("step_processor.renderer")

# Camera direction vectors: where the camera is PLACED relative to model center.
# Camera always looks toward center from direction * distance.
#
# Trimesh world axes after GLTF load (STEP Y and Z are swapped by GLTF Y-up conv):
#   X = width  — left/right
#   Y = depth  — Y_min = screen face (LCD), Y_max = back panel (ports)
#   Z = height — Z_min = TOP end of display, Z_max = BOTTOM end of display
#
# "front": camera at -Y (screen side) looks toward +Y → sees LCD face.
# "rear":  camera at +Y (back side) looks toward -Y  → sees port panel.
# "left":  camera at -X looks toward +X               → sees left edge.
# "right": camera at +X looks toward -X               → sees right edge.
# iso views: -Y component keeps camera on screen side; -Z = toward top end.
VIEW_CAMERAS = {
    "front":     ( 0, -1,  0),    # LCD screen face  (Y_min side)
    "rear":      ( 0,  1,  0),    # back panel/ports (Y_max side)
    "left":      (-1,  0,  0),    # left edge        (X_min side)
    "right":     ( 1,  0,  0),    # right edge       (X_max side)
    "top":       ( 0,  0, -1),    # top end          (Z_min side)
    "bottom":    ( 0,  0,  1),    # bottom end       (Z_max side)
    "iso_left":  (-1, -1, -0.5),  # front-left-above: screen + left edge + top
    "iso_right": ( 1, -1, -0.5),  # front-right-above: screen + right edge + top
}
DEFAULT_VIEWS  = ["front", "rear", "left", "right", "iso_left", "iso_right"]

# Color palette for per-part coloring when GLTF has no embedded colors
# 20 distinct RGBA colors (alpha=200 for slight transparency on overlaps)
PART_COLORS = [
    [180, 180, 185, 255],  # light steel
    [ 70, 130, 180, 255],  # steel blue
    [205, 133,  63, 255],  # peru / bronze
    [ 60, 179, 113, 255],  # medium sea green
    [188, 143, 143, 255],  # rosy brown
    [100, 149, 237, 255],  # cornflower blue
    [255, 160,  50, 255],  # dark orange
    [147, 112, 219, 255],  # medium purple
    [ 46, 139,  87, 255],  # sea green
    [205,  92,  92, 255],  # indian red
    [135, 206, 235, 255],  # sky blue
    [244, 164,  96, 255],  # sandy brown
    [106,  90, 205, 255],  # slate blue
    [ 32, 178, 170, 255],  # light sea green
    [220,  20,  60, 255],  # crimson
    [218, 165,  32, 255],  # goldenrod
    [ 72,  61, 139, 255],  # dark slate blue
    [143, 188, 143, 255],  # dark sea green
    [255,  99,  71, 255],  # tomato
    [176, 196, 222, 255],  # light steel blue
]
DEFAULT_WIDTH  = 1024
DEFAULT_HEIGHT = 768


def render_views(model: StepModel, step_path: Path,
                 views=None, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT) -> dict:
    """Render PNG views. Returns dict of view_name → output Path."""
    views = views or DEFAULT_VIEWS
    stem = step_path.stem
    out_dir = step_path.parent
    results = {}
    mesh = _get_mesh(model, step_path)
    if mesh is None:
        logger.warning("Could not obtain mesh — thumbnails skipped")
        return results
    for view_name in views:
        if view_name not in VIEW_CAMERAS:
            logger.warning(f"Unknown view '{view_name}' — skipping")
            continue
        out_path = out_dir / f"{stem}_{view_name}.png"
        try:
            _render_single_view(mesh, VIEW_CAMERAS[view_name], out_path, width, height)
            results[view_name] = out_path
            logger.info(f"Rendered: {out_path.name}")
        except Exception as e:
            logger.warning(f"Render failed for '{view_name}': {e}")
    return results


def _get_mesh(model: StepModel, step_path: Path):
    """Get a single assembled trimesh.Trimesh from the model.

    Always returns a single concatenated mesh (not a Scene) so the camera
    distance and bounds calculations are correct and parts don't explode.

    Priority:
      1. build123d → GLTF → scene.dump() → concatenate with transforms applied
         (preserves per-part colors if embedded in STEP)
      2. build123d → single STL of full assembly (fallback, monochrome but correct)
      3. FreeCAD path → bounding box box mesh
    """
    try:
        import trimesh
    except ImportError:
        logger.warning("trimesh not installed — thumbnails unavailable. pip install trimesh")
        return None

    if model.backend == "build123d":
        # ── GLTF path: color-aware, transforms flattened via scene.dump() ─────
        with tempfile.NamedTemporaryFile(suffix=".gltf", delete=False) as tmp:
            gltf_path = Path(tmp.name)
        try:
            from build123d import export_gltf
            import warnings
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                export_gltf(model.shape, str(gltf_path))
            scene = trimesh.load(str(gltf_path))

            if isinstance(scene, trimesh.Scene) and scene.geometry:
                # scene.dump() applies the full transform graph → parts at correct positions
                dumped = scene.dump()
                if dumped:
                    # Pull baseColorFactor directly from PBR material per mesh.
                    # .to_color() loses this in some trimesh versions.
                    materialized = []
                    for m in dumped:
                        try:
                            bc = m.visual.material.baseColorFactor
                            if bc is not None:
                                m.visual = trimesh.visual.ColorVisuals(
                                    mesh=m, face_colors=np.array(bc, dtype=np.uint8))
                            else:
                                m.visual = trimesh.visual.ColorVisuals(
                                    mesh=m, face_colors=[185, 190, 195, 255])
                        except Exception:
                            try:
                                m.visual = m.visual.to_color()
                            except Exception:
                                pass
                        materialized.append(m)

                    mesh = trimesh.util.concatenate(materialized)
                    if mesh is not None and len(mesh.faces) > 0:
                        n_colors = len(set(
                            tuple(c) for c in mesh.visual.face_colors[:, :3][::1000]
                        )) if hasattr(mesh.visual, 'face_colors') else 0
                        logger.info(f"GLTF: assembled {len(mesh.faces)} faces, "
                                    f"~{n_colors} distinct colors sampled")
                        return mesh
        except Exception as e:
            logger.warning(f"GLTF path failed ({e}) — falling back to STL")
        finally:
            try:
                gltf_path.unlink()
            except Exception:
                pass

        # ── STL fallback: single merged mesh, correct positions, monochrome ───
        with tempfile.NamedTemporaryFile(suffix=".stl", delete=False) as tmp:
            stl_path = Path(tmp.name)
        try:
            from build123d import export_stl
            export_stl(model.shape, str(stl_path))
            mesh = trimesh.load(str(stl_path), force="mesh")
            if mesh is not None and len(mesh.faces) > 0:
                mesh.visual.face_colors = [185, 190, 195, 255]
                logger.info(f"STL fallback: {len(mesh.faces)} faces (uniform color)")
                return mesh
        except Exception as e:
            logger.warning(f"STL fallback failed ({e})")
        finally:
            try:
                stl_path.unlink()
            except Exception:
                pass

    # FreeCAD / last resort: bounding box
    return _bbox_wireframe_mesh(model)


def _mesh_has_color(mesh) -> bool:
    """Return True if the mesh has meaningful (non-gray) face colors."""
    try:
        fc = mesh.visual.face_colors
        if fc is None or len(fc) == 0:
            return False
        # If all faces are within ±25 of [128,128,128] treat as uncolored
        gray = [128, 128, 128]
        mean_rgb = fc[:, :3].mean(axis=0)
        return not all(abs(int(mean_rgb[i]) - gray[i]) < 25 for i in range(3))
    except Exception:
        return False


def _bbox_wireframe_mesh(model: StepModel):
    """Create a simple box mesh from the model's bounding box. Last-resort fallback."""
    try:
        import trimesh
        bb = model.shape.bounding_box()
        # _FreeCADShapeProxy stores a _BoundBox
        if hasattr(bb, "XMin"):
            extents = [
                bb.XMax - bb.XMin,
                bb.YMax - bb.YMin,
                bb.ZMax - bb.ZMin,
            ]
            center = [
                (bb.XMax + bb.XMin) / 2,
                (bb.YMax + bb.YMin) / 2,
                (bb.ZMax + bb.ZMin) / 2,
            ]
        else:
            # build123d BoundBox
            extents = [bb.size.X, bb.size.Y, bb.size.Z]
            center = [bb.center.X, bb.center.Y, bb.center.Z]
        mesh = trimesh.creation.box(extents=extents)
        mesh.apply_translation(center)
        logger.debug("Using bbox wireframe mesh for rendering")
        return mesh
    except Exception as e:
        logger.warning(f"Bbox wireframe mesh failed: {e}")
        return None


def _render_single_view(mesh, camera_direction: tuple, out_path: Path,
                         width: int, height: int):
    """Render one view using pyrender offscreen, save to out_path."""
    try:
        import pyrender
    except ImportError:
        raise ImportError("pyrender not installed. pip install pyrender")

    # Normalize camera direction
    direction = np.array(camera_direction, dtype=float)
    direction = direction / np.linalg.norm(direction)

    # Bounding box of the assembled mesh
    bounds = mesh.bounds  # shape (2,3): [[xmin,ymin,zmin],[xmax,ymax,zmax]]
    center = (bounds[0] + bounds[1]) / 2.0
    diag = np.linalg.norm(bounds[1] - bounds[0])

    # Camera sits at 2.5× diagonal distance from center, looking at center
    camera_distance = diag * 2.5
    eye = center + direction * camera_distance

    # World up: (0,0,-1) = negative Z = top of display in trimesh GLTF space.
    # Fallback to (0,-1,0) for top/bottom end views where direction ≈ ±Z.
    world_up = np.array([0, 0, -1], dtype=float)
    if abs(np.dot(direction, world_up)) > 0.9:
        world_up = np.array([0, -1, 0], dtype=float)

    # Proven camera frame formula (right-handed, same structure as original code):
    #   right   = cross(world_up, direction)   [world_up × backward]
    #   cam_up  = cross(direction, right)       [backward × right]
    #   col2    = direction                     [camera +Z = backward; looks down -Z]
    right = np.cross(world_up, direction)
    if np.linalg.norm(right) < 1e-8:
        right = np.cross(np.array([1, 0, 0], dtype=float), direction)
    right = right / np.linalg.norm(right)
    cam_up = np.cross(direction, right)
    cam_up = cam_up / np.linalg.norm(cam_up)

    # 4×4 camera pose: columns = [right, cam_up, backward, eye]
    camera_pose = np.eye(4)
    camera_pose[:3, 0] = right
    camera_pose[:3, 1] = cam_up
    camera_pose[:3, 2] = direction   # camera +Z = backward; pyrender looks down -Z
    camera_pose[:3, 3] = eye

    # Build pyrender scene — white background
    pr_scene = pyrender.Scene(ambient_light=[0.35, 0.35, 0.35],
                              bg_color=[255, 255, 255, 255])
    pr_scene.add(pyrender.Mesh.from_trimesh(mesh, smooth=False))

    # FOV sized so the model fills ~80% of the frame
    yfov = 2.0 * np.arctan((diag * 0.5) / camera_distance) * 1.25
    camera = pyrender.PerspectiveCamera(yfov=yfov, aspectRatio=width / height)
    pr_scene.add(camera, pose=camera_pose)

    # Key light from camera position + fill from above-opposite
    pr_scene.add(pyrender.DirectionalLight(color=np.ones(3), intensity=4.5),
                 pose=camera_pose)
    fill_pose = np.eye(4)
    # Fill light: offset from top of display (-Z) to avoid zero-vector when
    # direction is parallel to (0,1,0) (e.g. rear view).
    fill_dir = -direction + np.array([0, 0, -1], dtype=float)
    fill_norm = np.linalg.norm(fill_dir)
    if fill_norm < 1e-8:
        fill_dir = np.array([0, 0, -1], dtype=float)
    else:
        fill_dir = fill_dir / fill_norm
    fill_eye = center - fill_dir * camera_distance
    fill_pose[:3, 3] = fill_eye
    pr_scene.add(pyrender.DirectionalLight(color=np.ones(3), intensity=1.8),
                 pose=fill_pose)

    # Offscreen render
    r = pyrender.OffscreenRenderer(viewport_width=width, viewport_height=height)
    try:
        color, _ = r.render(pr_scene)
    finally:
        r.delete()

    from PIL import Image
    Image.fromarray(color).save(str(out_path))