services/mesh_utils.py

"""
services/mesh_utils.py
-----------------------
Tumor mesh generation from a segmentation mask.

This is a PLACEHOLDER implementation. It consumes the M1 output mask
(pred_mask.npy — a (D, H, W) int label volume) and produces a triangle
mesh via marching cubes, then writes it as .obj (always) and .glb
(best-effort, via trimesh if available).

It deliberately does NOT touch the MRI volume itself — only the mask —
per the pipeline contract.

Swap-out point:
    Replace `generate_tumor_mesh()` internals with a real meshing
    pipeline (e.g. smoothed marching cubes + decimation + a proper
    glTF exporter) without changing its signature, and the rest of
    pipeline.py keeps working unmodified.
"""

from __future__ import annotations

import os
from dataclasses import dataclass
from typing import Optional

import numpy as np
from skimage import measure


@dataclass
class MeshResult:
    obj_path: str
    glb_path: Optional[str]
    n_vertices: int
    n_faces: int
    label_value: int
    error: Optional[str] = None

    def to_dict(self) -> dict:
        return {
            "obj_path": self.obj_path,
            "glb_path": self.glb_path,
            "n_vertices": self.n_vertices,
            "n_faces": self.n_faces,
            "label_value": self.label_value,
            "error": self.error,
        }


def _write_obj(path: str, verts: np.ndarray, faces: np.ndarray) -> None:
    with open(path, "w") as f:
        for v in verts:
            f.write(f"v {v[0]:.6f} {v[1]:.6f} {v[2]:.6f}\n")
        for face in faces:
            # OBJ is 1-indexed
            f.write(f"f {face[0] + 1} {face[1] + 1} {face[2] + 1}\n")


def _write_glb(path: str, verts: np.ndarray, faces: np.ndarray) -> bool:
    """
    Best-effort .glb export via trimesh. Returns True on success,
    False if trimesh isn't installed or export fails — callers should
    treat a missing .glb as non-fatal (the .obj is the source of truth).
    """
    try:
        import trimesh  # optional dependency

        mesh = trimesh.Trimesh(vertices=verts, faces=faces, process=False)
        mesh.export(path)
        return True
    except Exception:
        return False


def generate_tumor_mesh(
    mask_path: str,
    out_dir: str,
    label_value: int = 1,
    level: float = 0.5,
    step_size: int = 1,
) -> MeshResult:
    """
    Generate a tumor surface mesh from a saved segmentation mask.

    Args:
        mask_path:   path to pred_mask.npy ((D, H, W) int array, as
                     produced by M1's main()/postprocessing).
        out_dir:     directory to write tumor.obj / tumor.glb into.
        label_value: which label in the mask to mesh. BraTS-style masks
                     are multi-class (e.g. edema / enhancing / necrotic);
                     callers needing a "whole tumor" mesh should binarize
                     before calling this, or call once per label.
        level:       marching-cubes iso-level on the binarized mask.
        step_size:   marching-cubes step size (>1 downsamples for speed).

    Returns:
        MeshResult with paths and basic stats. On failure, obj_path/
        glb_path will be empty strings and `error` will be set —
        callers should treat meshing as best-effort and not abort the
        rest of the pipeline if it fails.
    """
    os.makedirs(out_dir, exist_ok=True)
    obj_path = os.path.join(out_dir, "tumor.obj")
    glb_path = os.path.join(out_dir, "tumor.glb")

    try:
        mask = np.load(mask_path)
        binary = (mask == label_value).astype(np.uint8)

        if binary.sum() == 0:
            return MeshResult(
                obj_path="", glb_path=None, n_vertices=0, n_faces=0,
                label_value=label_value,
                error=f"No voxels found for label_value={label_value}; mesh not generated.",
            )

        verts, faces, _normals, _values = measure.marching_cubes(
            binary, level=level, step_size=step_size
        )

        _write_obj(obj_path, verts, faces)
        glb_written = _write_glb(glb_path, verts, faces)

        return MeshResult(
            obj_path=obj_path,
            glb_path=glb_path if glb_written else None,
            n_vertices=int(verts.shape[0]),
            n_faces=int(faces.shape[0]),
            label_value=label_value,
        )

    except Exception as e:
        return MeshResult(
            obj_path="", glb_path=None, n_vertices=0, n_faces=0,
            label_value=label_value,
            error=f"Mesh generation failed: {e}",
        )