HartOMat/threejs-renderer/app.py

"""
Three.js renderer service — FastAPI microservice.

Pipeline: STEP → STL (cadquery) → Three.js in headless Chromium → PNG screenshot.

Two render modes
────────────────
  part_colors = None      Single grey metallic mesh (original behaviour).
  part_colors = dict      Connected-component analysis in JavaScript:
                          The STL is loaded as one mesh; disconnected islands are
                          detected entirely in the browser and each gets a distinct
                          palette colour.  No server-side OCC/per-part extraction —
                          just one STL conversion and client-side graph analysis.
"""
import asyncio
import base64
import json
import logging
from pathlib import Path

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

logger = logging.getLogger(__name__)

app = FastAPI(title="Three.js Renderer", version="1.0.0")

# 10-colour palette used for connected-component assignment
PALETTE = [
    "#4C9BE8", "#E85B4C", "#4CBE72", "#E8A84C", "#A04CE8",
    "#4CD4E8", "#E84CA8", "#7EC850", "#E86B30", "#5088C8",
]


class RenderRequest(BaseModel):
    step_path: str
    output_path: str
    width: int = 512
    height: int = 512
    # None  → single grey mesh
    # {}    → auto-colour by connected-component index (palette)
    # {...} → same (named-part colour mapping is handled in JS if names match)
    part_colors: dict[str, str] | None = None
    rotation_x: float = 0.0
    rotation_y: float = 0.0
    rotation_z: float = 0.0


@app.get("/health")
async def health():
    return {"status": "ok", "renderer": "threejs"}


@app.post("/render")
async def render(req: RenderRequest):
    step_path = Path(req.step_path)
    output_path = Path(req.output_path)

    if not step_path.exists():
        raise HTTPException(404, detail=f"STEP file not found: {step_path}")

    output_path.parent.mkdir(parents=True, exist_ok=True)

    # Persistent STL cache — same convention as blender-renderer (quality always "low")
    stl_path = step_path.parent / f"{step_path.stem}_low.stl"
    if not stl_path.exists() or stl_path.stat().st_size == 0:
        try:
            _convert_step_to_stl(step_path, stl_path)
        except Exception as e:
            logger.error(f"STEP→STL conversion failed: {e}")
            raise HTTPException(500, detail=f"STEP conversion failed: {e}")
        logger.info("STL cached: %s (%d KB)", stl_path.name, stl_path.stat().st_size // 1024)
    else:
        logger.info("STL cache hit: %s (%d KB)", stl_path.name, stl_path.stat().st_size // 1024)

    use_colors = req.part_colors is not None
    try:
        await asyncio.to_thread(
            _render_stl_threejs,
            stl_path, output_path, req.width, req.height, use_colors,
            req.rotation_x, req.rotation_y, req.rotation_z,
        )
    except Exception as e:
        logger.error(f"Three.js render failed: {e}")
        raise HTTPException(500, detail=f"Three.js render failed: {e}")

    if not output_path.exists():
        raise HTTPException(500, detail="Render produced no output file")

    return {
        "output_path": str(output_path),
        "status": "ok",
        "renderer": "threejs-colored" if use_colors else "threejs",
    }


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _convert_step_to_stl(step_path: Path, stl_path: Path) -> None:
    """Convert STEP to a single binary STL via cadquery."""
    import cadquery as cq
    shape = cq.importers.importStep(str(step_path))
    cq.exporters.export(shape, str(stl_path))
    if not stl_path.exists() or stl_path.stat().st_size == 0:
        raise RuntimeError("cadquery produced empty STL")


def _render_stl_threejs(
    stl_path: Path,
    output_path: Path,
    width: int,
    height: int,
    use_colors: bool,
    rotation_x: float = 0.0,
    rotation_y: float = 0.0,
    rotation_z: float = 0.0,
) -> None:
    """Render STL via Three.js in headless Chromium."""
    from playwright.sync_api import sync_playwright

    stl_b64 = base64.b64encode(stl_path.read_bytes()).decode()
    filename = stl_path.stem
    palette_json = json.dumps(PALETTE)

    html = _build_html(stl_b64, filename, width, height, palette_json, use_colors,
                       rotation_x, rotation_y, rotation_z)

    with sync_playwright() as p:
        browser = p.chromium.launch(
            args=["--no-sandbox", "--disable-setuid-sandbox", "--disable-gpu"]
        )
        page = browser.new_page(viewport={"width": width, "height": height})
        page.set_content(html, wait_until="domcontentloaded")
        try:
            page.wait_for_function("window.__renderDone === true", timeout=12000)
        except Exception:
            pass   # take screenshot anyway
        page.screenshot(
            path=str(output_path),
            full_page=False,
            clip={"x": 0, "y": 0, "width": width, "height": height},
        )
        browser.close()


def _build_html(
    stl_b64: str,
    filename: str,
    width: int,
    height: int,
    palette_json: str,
    use_colors: bool,
    rotation_x: float = 0.0,
    rotation_y: float = 0.0,
    rotation_z: float = 0.0,
) -> str:
    """
    Build a self-contained HTML page that renders the STL with Three.js.

    When use_colors=True the JavaScript runs a Union-Find connected-component
    analysis on the vertex graph of the STL and paints each disconnected island
    with a distinct colour from the palette.  This requires no server-side part
    extraction — it works directly on the flat triangle soup in the STL.
    """

    # ---------- colour-assignment script (injected only when use_colors=True)
    color_script = ""
    if use_colors:
        color_script = f"""
// ── Connected-component colouring ─────────────────────────────────────────
// Each STL face is a triplet of un-shared vertices.  We weld coincident
// vertices by their rounded position string, then run Union-Find on the
// resulting graph to identify disconnected parts.  Each part gets a colour
// from the palette.
function applyPartColors(geometry, palette) {{
  const pos = geometry.attributes.position;
  const n   = pos.count;

  // Round to 4 d.p. to merge floating-point near-duplicates
  const key = i =>
    Math.round(pos.getX(i)*1e4) + ',' +
    Math.round(pos.getY(i)*1e4) + ',' +
    Math.round(pos.getZ(i)*1e4);

  // Map position string → canonical vertex index
  const posMap = Object.create(null);
  const canon  = new Int32Array(n);
  for (let i = 0; i < n; i++) {{
    const k = key(i);
    if (posMap[k] === undefined) posMap[k] = i;
    canon[i] = posMap[k];
  }}

  // Union-Find with path compression + union by rank
  const parent = new Int32Array(n);
  const rank   = new Uint8Array(n);
  for (let i = 0; i < n; i++) parent[i] = i;
  function find(x) {{
    while (parent[x] !== x) {{ parent[x] = parent[parent[x]]; x = parent[x]; }}
    return x;
  }}
  function unite(a, b) {{
    a = find(a); b = find(b);
    if (a === b) return;
    if (rank[a] < rank[b]) {{ let t = a; a = b; b = t; }}
    parent[b] = a;
    if (rank[a] === rank[b]) rank[a]++;
  }}

  // Connect the three canonical vertices of every triangle
  for (let i = 0; i < n; i += 3) {{
    unite(canon[i], canon[i+1]);
    unite(canon[i+1], canon[i+2]);
  }}

  // Assign a palette index to each component root
  const compIdx = Object.create(null);
  let nextIdx = 0;
  const colors = new Float32Array(n * 3);
  for (let i = 0; i < n; i++) {{
    const root = find(canon[i]);
    if (compIdx[root] === undefined) compIdx[root] = nextIdx++;
    const hex = palette[compIdx[root] % palette.length];
    colors[i*3]   = parseInt(hex.slice(1,3), 16) / 255;
    colors[i*3+1] = parseInt(hex.slice(3,5), 16) / 255;
    colors[i*3+2] = parseInt(hex.slice(5,7), 16) / 255;
  }}

  geometry.setAttribute('color', new THREE.BufferAttribute(colors, 3));
  return new THREE.MeshStandardMaterial({{
    vertexColors: true,
    metalness: 0.4,
    roughness: 0.45,
  }});
}}

const palette = {palette_json};
const material = applyPartColors(geometry, palette);
"""
    else:
        color_script = """
const material = new THREE.MeshStandardMaterial({
  color: 0xc0cad8,
  metalness: 0.8,
  roughness: 0.3,
});
"""

    return f"""<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<style>
  * {{ margin:0; padding:0; box-sizing:border-box; }}
  body {{ background:#f5f6f8; overflow:hidden; }}
  canvas {{ display:block; }}
</style>
</head>
<body>
<script type="importmap">
{{
  "imports": {{
    "three": "https://cdn.jsdelivr.net/npm/three@0.162.0/build/three.module.js",
    "three/addons/": "https://cdn.jsdelivr.net/npm/three@0.162.0/examples/jsm/"
  }}
}}
</script>
<script type="module">
import * as THREE from 'three';
import {{ STLLoader }} from 'three/addons/loaders/STLLoader.js';

const W = {width}, H = {height};
const renderer = new THREE.WebGLRenderer({{ antialias: true }});
renderer.setSize(W, H);
renderer.setPixelRatio(1);
renderer.outputColorSpace = THREE.SRGBColorSpace;
renderer.shadowMap.enabled = true;
renderer.shadowMap.type = THREE.PCFSoftShadowMap;
document.body.appendChild(renderer.domElement);

const scene = new THREE.Scene();
scene.background = new THREE.Color(0xf5f6f8);

const camera = new THREE.PerspectiveCamera(45, W / H, 0.001, 10000);
scene.add(camera);

const ambientLight = new THREE.AmbientLight(0xffffff, 0.6);
scene.add(ambientLight);
const dirLight = new THREE.DirectionalLight(0xffffff, 2.5);
dirLight.position.set(1, 2, 1.5);
dirLight.castShadow = true;
scene.add(dirLight);
const fillLight = new THREE.DirectionalLight(0xddeeff, 1.0);
fillLight.position.set(-1, -0.5, -1);
scene.add(fillLight);

// Decode base64 STL
const b64    = "{stl_b64}";
const binary = Uint8Array.from(atob(b64), c => c.charCodeAt(0));
const loader = new STLLoader();
const geometry = loader.parse(binary.buffer);
geometry.computeVertexNormals();

// Material (grey or per-part coloured)
{color_script}

const mesh = new THREE.Mesh(geometry, material);
mesh.castShadow = true;
mesh.receiveShadow = true;
scene.add(mesh);

// Centre and frame
geometry.computeBoundingBox();
const box = geometry.boundingBox;
const center = new THREE.Vector3();
box.getCenter(center);
mesh.position.sub(center);
mesh.rotation.set({rotation_x}*Math.PI/180, {rotation_y}*Math.PI/180, {rotation_z}*Math.PI/180);

const size = new THREE.Vector3();
box.getSize(size);
const maxDim = Math.max(size.x, size.y, size.z);
const fov = camera.fov * (Math.PI / 180);
let dist = maxDim / (2 * Math.tan(fov / 2)) * 1.15;
dist = Math.max(dist, 0.01);
camera.position.set(dist * 0.8, dist * 0.6, dist * 0.8);
camera.lookAt(0, 0, 0);

// Schaeffler green top bar
const topBar = document.createElement('div');
topBar.style.cssText =
  `position:fixed;top:0;left:0;width:${{W}}px;height:10px;background:#00893d;z-index:10`;
document.body.appendChild(topBar);

// Model name label
const label = document.createElement('div');
label.textContent = "{filename}";
label.style.cssText =
  `position:fixed;bottom:0;left:0;width:${{W}}px;background:rgba(20,35,55,0.85);` +
  `color:#fff;text-align:center;font-size:13px;font-family:monospace;` +
  `padding:6px 4px;box-sizing:border-box;z-index:10`;
document.body.appendChild(label);

renderer.render(scene, camera);
window.__renderDone = true;
</script>
</body>
</html>"""