Skip to content

feat(mesh): add unified mesh operations registry #329

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
DrHepa wants to merge 1 commit into
base: dev
Choose a base branch
from
Open

feat(mesh): add unified mesh operations registry #329

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
285 changes: 101 additions & 184 deletions api/routers/optimize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,27 +1,24 @@
import hashlib
import os
import re
import shutil
import tempfile
import uuid

try:
import pymeshlab as _pymeshlab
_PYMESHLAB_AVAILABLE = True
except ImportError:
_pymeshlab = None
_PYMESHLAB_AVAILABLE = False

import numpy as np
import trimesh
import trimesh.visual
from fastapi import APIRouter, HTTPException, UploadFile, File
from fastapi.responses import FileResponse, Response
from pathlib import Path
from urllib.parse import quote
from pydantic import BaseModel
from pydantic import BaseModel, Field

from services.generator_registry import WORKSPACE_DIR
from services.mesh_ops import (
MeshOpContext,
MeshOpNotFoundError,
MeshOpResult,
MeshOpUnavailableError,
mesh_ops_registry,
)

router = APIRouter(tags=["optimize"])

Expand All @@ -36,16 +33,16 @@ class SmoothRequest(BaseModel):
iterations: int


class MeshOpRequest(BaseModel):
path: str
params: dict[str, object] = Field(default_factory=dict)


class TransformRequest(BaseModel):
path: str # format: "{collection}/{filename}"
matrix: list[list[float]] # row-major 4x4 world transform


def _require_pymeshlab():
if not _PYMESHLAB_AVAILABLE:
raise HTTPException(503, "pymeshlab is unavailable on this system (DLL blocked by Windows Application Control policy)")


def _resolve_input_path(raw_path: str) -> Path:
candidate = Path(raw_path)
if candidate.is_absolute():
Expand All @@ -62,142 +59,109 @@ def _resolve_input_path(raw_path: str) -> Path:
return resolved


@router.post("/mesh")
def optimize_mesh(body: OptimizeRequest):
_require_pymeshlab()
target_faces = max(100, min(500_000, body.target_faces))
def _operation_output_path(input_path: Path, output_name: str) -> Path:
workspace = WORKSPACE_DIR.resolve()
resolved_input = input_path.resolve()
output_dir = (
input_path.parent
if resolved_input == workspace or workspace in resolved_input.parents
else WORKSPACE_DIR / "Workflows"
)
output_dir.mkdir(parents=True, exist_ok=True)
return output_dir / output_name


def _run_operation(
operation_id: str,
input_path: Path,
params: dict[str, object],
output_path: Path | None = None,
preserve_visuals: bool = False,
) -> MeshOpResult:
context = MeshOpContext(
workspace_dir=WORKSPACE_DIR,
temp_dir=Path(tempfile.gettempdir()),
output_path=output_path,
preserve_visuals=preserve_visuals,
)
try:
return mesh_ops_registry.run(operation_id, input_path, params, context)
except MeshOpNotFoundError as exc:
raise HTTPException(404, f"Unknown mesh operation: {operation_id}") from exc
except MeshOpUnavailableError as exc:
raise HTTPException(503, str(exc)) from exc
except (TypeError, ValueError) as exc:
raise HTTPException(400, str(exc)) from exc

input_path = _resolve_input_path(body.path)

tmp_dir = tempfile.mkdtemp()
def _operation_response(result: MeshOpResult) -> dict[str, object]:
output_path = result.file_path.resolve()
try:
result = _decimate(str(input_path), target_faces, tmp_dir)
finally:
shutil.rmtree(tmp_dir, ignore_errors=True)

stem = input_path.stem
output_name = f"{stem}_opt{target_faces}.glb"
output_dir = input_path.parent if str(input_path).startswith(str(WORKSPACE_DIR.resolve())) else WORKSPACE_DIR / "Workflows"
output_dir.mkdir(parents=True, exist_ok=True)
output_path = output_dir / output_name
result.export(str(output_path))
relative_path = output_path.relative_to(WORKSPACE_DIR.resolve()).as_posix()
except ValueError:
payload: dict[str, object] = {"path": str(output_path)}
else:
payload = {
"path": relative_path,
"url": f"/workspace/{relative_path}",
}
payload.update(result.details)
return payload

face_count = len(result.faces)
rel = output_path.relative_to(WORKSPACE_DIR).as_posix()
return {"url": f"/workspace/{rel}", "face_count": face_count}

@router.get("/ops")
def list_mesh_operations():
return mesh_ops_registry.describe()

def _has_texture(geom: trimesh.Trimesh) -> bool:
if not isinstance(geom.visual, trimesh.visual.TextureVisuals):
return False
mat = geom.visual.material
if mat is None:
return False
# Simple material (SimpleMaterial / Material)
if getattr(mat, "image", None) is not None:
return True
# PBR material (from Trellis2 SLaT texturing and GLB imports)
if getattr(mat, "baseColorTexture", None) is not None:
return True
return False


def _get_texture_image(geom: trimesh.Trimesh):
"""Return the base color texture image regardless of material type."""
mat = geom.visual.material
img = getattr(mat, "image", None)
if img is not None:
return img
return getattr(mat, "baseColorTexture", None)


def _decimate(input_path: str, target_faces: int, tmp_dir: str) -> trimesh.Trimesh:
loaded = trimesh.load(input_path)
if isinstance(loaded, trimesh.Scene):
geoms = list(loaded.geometry.values())
geom = trimesh.util.concatenate(geoms) if len(geoms) > 1 else geoms[0]
else:
geom = loaded

ms = _pymeshlab.MeshSet()

if _has_texture(geom):
# ── Textured path: OBJ intermediate to preserve UV coordinates ──────
obj_in = os.path.join(tmp_dir, "input.obj")
mtl_in = os.path.join(tmp_dir, "input.mtl")
tex_in = os.path.join(tmp_dir, "texture.png")
obj_out = os.path.join(tmp_dir, "output.obj")

# Save texture image under a known filename (handles PBR and simple materials)
_get_texture_image(geom).save(tex_in)

# Export OBJ (trimesh writes UV coords + MTL)
geom.export(obj_in)

# Patch MTL so any map_Kd points to our known texture filename
if os.path.exists(mtl_in):
mtl = open(mtl_in).read()
mtl = re.sub(r"map_Kd\s+\S+", "map_Kd texture.png", mtl)
open(mtl_in, "w").write(mtl)

ms.load_new_mesh(obj_in)
ms.meshing_decimation_quadric_edge_collapse(
targetfacenum=target_faces,
preservetexcoord=True, # ← keeps UV coordinates intact
preservenormal=True,
preservetopology=True,
autoclean=True,
)
ms.save_current_mesh(obj_out)

# Patch output MTL too, so trimesh can find the texture on load
mtl_out = obj_out.replace(".obj", ".mtl")
if os.path.exists(mtl_out):
mtl = open(mtl_out).read()
mtl = re.sub(r"map_Kd\s+\S+", "map_Kd texture.png", mtl)
open(mtl_out, "w").write(mtl)
@router.post("/op/{op_name}")
def run_mesh_operation(op_name: str, body: MeshOpRequest):
input_path = _resolve_input_path(body.path)
return _operation_response(
_run_operation(op_name, input_path, body.params)
)

return trimesh.load(obj_out)

else:
# ── Geometry-only path: PLY (fast, no texture to worry about) ────────
ply_in = os.path.join(tmp_dir, "input.ply")
ply_out = os.path.join(tmp_dir, "output.ply")

geom.export(ply_in)
ms.load_new_mesh(ply_in)
ms.meshing_decimation_quadric_edge_collapse(
targetfacenum=target_faces,
preservenormal=True,
preservetopology=True,
autoclean=True,
@router.post("/mesh")
def optimize_mesh(body: OptimizeRequest):
target_faces = max(100, min(500_000, body.target_faces))
input_path = _resolve_input_path(body.path)
output_path = _operation_output_path(
input_path,
f"{input_path.stem}_opt{target_faces}.glb",
)
response = _operation_response(
_run_operation(
"decimate",
input_path,
{"target_faces": target_faces},
output_path,
)
ms.save_current_mesh(ply_out)
return trimesh.load(ply_out, force="mesh")
)
return {
"url": response.get("url"),
"face_count": response.get("face_count", 0),
}


@router.post("/smooth")
def smooth_mesh(body: SmoothRequest):
_require_pymeshlab()
iterations = max(1, min(20, body.iterations))

input_path = _resolve_input_path(body.path)

tmp_dir = tempfile.mkdtemp()
try:
result = _smooth(str(input_path), iterations, tmp_dir)
finally:
shutil.rmtree(tmp_dir, ignore_errors=True)

stem = input_path.stem
output_name = f"{stem}_smooth{iterations}.glb"
output_dir = input_path.parent if str(input_path).startswith(str(WORKSPACE_DIR.resolve())) else WORKSPACE_DIR / "Workflows"
output_dir.mkdir(parents=True, exist_ok=True)
output_path = output_dir / output_name
result.export(str(output_path))

rel = output_path.relative_to(WORKSPACE_DIR).as_posix()
return {"url": f"/workspace/{rel}"}
output_path = _operation_output_path(
input_path,
f"{input_path.stem}_smooth{iterations}.glb",
)
response = _operation_response(
_run_operation(
"smooth",
input_path,
{"iterations": iterations, "lambda_": 0.5, "mode": "laplacian"},
output_path,
preserve_visuals=True,
)
)
return {"url": response.get("url")}


@router.post("/transform")
Expand Down Expand Up @@ -228,53 +192,6 @@ def transform_mesh(body: TransformRequest):
return {"url": f"/workspace/{rel}"}


def _smooth(input_path: str, iterations: int, tmp_dir: str) -> trimesh.Trimesh:
loaded = trimesh.load(input_path)
if isinstance(loaded, trimesh.Scene):
geoms = list(loaded.geometry.values())
geom = trimesh.util.concatenate(geoms) if len(geoms) > 1 else geoms[0]
else:
geom = loaded

ms = _pymeshlab.MeshSet()

if _has_texture(geom):
obj_in = os.path.join(tmp_dir, "input.obj")
mtl_in = os.path.join(tmp_dir, "input.mtl")
tex_in = os.path.join(tmp_dir, "texture.png")
obj_out = os.path.join(tmp_dir, "output.obj")

_get_texture_image(geom).save(tex_in)
geom.export(obj_in)

if os.path.exists(mtl_in):
mtl = open(mtl_in).read()
mtl = re.sub(r"map_Kd\s+\S+", "map_Kd texture.png", mtl)
open(mtl_in, "w").write(mtl)

ms.load_new_mesh(obj_in)
ms.apply_coord_laplacian_smoothing(stepsmoothnum=iterations)
ms.save_current_mesh(obj_out)

mtl_out = obj_out.replace(".obj", ".mtl")
if os.path.exists(mtl_out):
mtl = open(mtl_out).read()
mtl = re.sub(r"map_Kd\s+\S+", "map_Kd texture.png", mtl)
open(mtl_out, "w").write(mtl)

return trimesh.load(obj_out)

else:
ply_in = os.path.join(tmp_dir, "input.ply")
ply_out = os.path.join(tmp_dir, "output.ply")

geom.export(ply_in)
ms.load_new_mesh(ply_in)
ms.apply_coord_laplacian_smoothing(stepsmoothnum=iterations)
ms.save_current_mesh(ply_out)
return trimesh.load(ply_out, force="mesh")


class ImportByPathRequest(BaseModel):
path: str # absolute path on disk

Expand Down Expand Up @@ -503,4 +420,4 @@ def export_mesh(path: str, format: str):
content=data,
media_type=mime,
headers={"Content-Disposition": f'attachment; filename="{stem}.{format}"'},
)
)
26 changes: 26 additions & 0 deletions api/services/mesh_ops/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,26 @@
"""Unified mesh operation registry used by the API and workflow nodes."""

from .builtin import BUILTIN_MESH_OPS
from .registry import MeshOpsRegistry
from .types import (
MeshOp,
MeshOpContext,
MeshOpExecutionError,
MeshOpNotFoundError,
MeshOpResult,
MeshOpUnavailableError,
)


mesh_ops_registry = MeshOpsRegistry(BUILTIN_MESH_OPS)

__all__ = [
"MeshOp",
"MeshOpContext",
"MeshOpExecutionError",
"MeshOpNotFoundError",
"MeshOpResult",
"MeshOpsRegistry",
"MeshOpUnavailableError",
"mesh_ops_registry",
]
Loading