app.py

"""Gradio demo for the complete HiTOPS single-mesh pipeline."""

from __future__ import annotations

import os
import shutil
import tempfile
import time
import uuid
from pathlib import Path

import gradio as gr
import numpy as np
import trimesh

from scripts.run_batch_sqfit import process_one_isolated


PROJECT_ROOT = Path(__file__).resolve().parent
JOB_ROOT = Path(tempfile.gettempdir()) / "hitops-gradio"
ALLOWED_SUFFIXES = {".ply", ".obj", ".stl", ".glb", ".gltf"}

MAX_UPLOAD_MB = int(os.getenv("HITOPS_MAX_UPLOAD_MB", "50"))
MAX_FACES = int(os.getenv("HITOPS_MAX_FACES", "1000000"))
JOB_TIMEOUT_SEC = int(os.getenv("HITOPS_JOB_TIMEOUT_SEC", "1800"))
JOB_TTL_SEC = int(os.getenv("HITOPS_JOB_TTL_SEC", "21600"))


def _cleanup_stale_jobs() -> None:
    """Remove expired jobs from the dedicated temporary-job directory."""
    JOB_ROOT.mkdir(parents=True, exist_ok=True)
    cutoff = time.time() - JOB_TTL_SEC
    for path in JOB_ROOT.iterdir():
        try:
            if path.is_dir() and path.stat().st_mtime < cutoff:
                shutil.rmtree(path)
        except OSError:
            # Another worker may still own or already have removed the path.
            continue


def _load_and_validate_mesh(uploaded_mesh: str) -> trimesh.Trimesh:
    """Load an uploaded mesh and enforce the public-demo resource contract."""
    source = Path(uploaded_mesh)
    if not source.is_file():
        raise gr.Error("The uploaded mesh file is unavailable.")
    if source.suffix.lower() not in ALLOWED_SUFFIXES:
        supported = ", ".join(sorted(ALLOWED_SUFFIXES))
        raise gr.Error(f"Unsupported file type. Supported formats: {supported}.")

    size_mb = source.stat().st_size / (1024 * 1024)
    if size_mb > MAX_UPLOAD_MB:
        raise gr.Error(
            f"The upload is {size_mb:.1f} MB; the limit is {MAX_UPLOAD_MB} MB."
        )

    try:
        mesh = trimesh.load(source, force="mesh", process=False)
    except Exception as exc:
        raise gr.Error(f"Failed to read the mesh: {exc}") from exc

    if not isinstance(mesh, trimesh.Trimesh) or mesh.is_empty:
        raise gr.Error("The upload does not contain a valid triangle mesh.")
    if len(mesh.vertices) == 0 or len(mesh.faces) == 0:
        raise gr.Error("The mesh has no vertices or triangular faces.")
    if len(mesh.faces) > MAX_FACES:
        raise gr.Error(
            f"The mesh has {len(mesh.faces):,} faces; the limit is "
            f"{MAX_FACES:,} faces."
        )
    if not np.isfinite(np.asarray(mesh.vertices)).all():
        raise gr.Error("The mesh contains NaN or infinite vertex coordinates.")
    if not mesh.is_watertight:
        gr.Warning(
            "The mesh is not watertight. HiTOPS will try to process it, but "
            "watertight preprocessing is recommended if the pipeline fails."
        )
    return mesh


def _export_glb(mesh: trimesh.Trimesh, target: Path) -> str:
    """Export a browser-friendly GLB while preserving mesh vertex colors."""
    target.parent.mkdir(parents=True, exist_ok=True)
    try:
        mesh.export(target, file_type="glb")
    except Exception as exc:
        raise gr.Error(f"Failed to create the browser preview: {exc}") from exc
    if not target.is_file() or target.stat().st_size == 0:
        raise gr.Error("The browser preview GLB was not created.")
    return str(target)


def prepare_input_preview(uploaded_mesh: str | None) -> str | None:
    """Validate an upload and convert it to GLB for reliable WebGL display."""
    if not uploaded_mesh:
        return None
    _cleanup_stale_jobs()
    mesh = _load_and_validate_mesh(uploaded_mesh)
    preview_dir = JOB_ROOT / f"preview-{uuid.uuid4().hex[:12]}"
    return _export_glb(mesh, preview_dir / "input-preview.glb")


def run_hitops(
    uploaded_mesh: str | None,
    progress: gr.Progress = gr.Progress(),
):
    """Run the complete SQ-fit, curvature-segmentation, and mapping pipeline."""
    if not uploaded_mesh:
        raise gr.Error("Upload a 3D triangle mesh first.")

    _cleanup_stale_jobs()
    progress(0.03, desc="Validating mesh")
    mesh = _load_and_validate_mesh(uploaded_mesh)

    job_id = uuid.uuid4().hex[:12]
    job_dir = JOB_ROOT / job_id
    mesh_root = job_dir / "inputs"
    mesh_dir = mesh_root / "00" / job_id
    output_root = job_dir / "outputs"
    mesh_dir.mkdir(parents=True, exist_ok=False)

    # The batch driver expects <mesh_root>/<shard>/<uid>/full.ply. Exporting
    # performs a real format conversion instead of merely renaming the upload.
    input_ply = mesh_dir / "full.ply"
    try:
        mesh.export(input_ply, file_type="ply")
    except Exception as exc:
        shutil.rmtree(job_dir, ignore_errors=True)
        raise gr.Error(f"Failed to convert the mesh to PLY: {exc}") from exc

    progress(0.08, desc="Running HiTOPS (this can take several minutes)")
    try:
        result = process_one_isolated(
            uid=job_id,
            mesh_root=str(mesh_root),
            shard="00",
            output_root=str(output_root),
            timeout_sec=JOB_TIMEOUT_SEC,
        )
    except Exception as exc:
        raise gr.Error(f"HiTOPS failed: {exc}") from exc

    if result.get("status") != "ok":
        error = result.get("error", "unknown pipeline error")
        raise gr.Error(f"HiTOPS failed: {error}")

    result_dir = output_root / job_id
    mapping_dir = result_dir / "mesh_mapping_v8"
    colored_mesh = mapping_dir / "mesh_mapped_v8.ply"
    if not colored_mesh.is_file():
        raise gr.Error("HiTOPS completed without producing the colored result mesh.")

    progress(0.94, desc="Creating browser preview")
    try:
        result_mesh = trimesh.load(colored_mesh, force="mesh", process=False)
    except Exception as exc:
        raise gr.Error(f"Failed to read the colored result mesh: {exc}") from exc
    preview_glb = _export_glb(result_mesh, mapping_dir / "mesh_mapped_v8.glb")

    progress(0.97, desc="Packaging results")
    archive_path = shutil.make_archive(
        str(job_dir / f"hitops-{job_id}-results"),
        "zip",
        root_dir=output_root,
        base_dir=job_id,
    )

    progress(1.0, desc="Complete")
    return preview_glb, result, archive_path


example_files = [
    [str(path)]
    for path in sorted((PROJECT_ROOT / "examples" / "hy3d" / "00").glob("*/full.ply"))
]

with gr.Blocks(title="HiTOPS") as demo:
    gr.Markdown(
        """
# HiTOPS

**Geometry-only 3D part decomposition with adaptive structural carriers and
superquadrics.**

Upload a triangle mesh to run SQ fitting, curvature segmentation, and final
face-to-part mapping. Watertight input is recommended. Processing may take several minutes.
"""
    )

    with gr.Row():
        with gr.Column():
            input_file = gr.File(
                label="Input mesh file",
                file_types=sorted(ALLOWED_SUFFIXES),
                type="filepath",
            )
            input_preview = gr.Model3D(
                label="Input preview",
                height=470,
                interactive=False,
                clear_color=(0.08, 0.08, 0.10, 1.0),
            )
        output_mesh = gr.Model3D(
            label="HiTOPS part decomposition",
            height=520,
            interactive=False,
            clear_color=(0.08, 0.08, 0.10, 1.0),
        )

    run_button = gr.Button("Run HiTOPS", variant="primary")

    with gr.Row():
        run_info = gr.JSON(label="Run summary")
        download = gr.File(label="Download complete results")

    input_file.change(
        fn=prepare_input_preview,
        inputs=input_file,
        outputs=input_preview,
        concurrency_limit=2,
        api_name=False,
    )

    run_button.click(
        fn=run_hitops,
        inputs=input_file,
        outputs=[output_mesh, run_info, download],
        concurrency_limit=1,
        api_name="run_hitops",
    )

    if example_files:
        gr.Examples(
            examples=example_files,
            inputs=input_file,
            label="Bundled HY3D examples",
            cache_examples=False,
        )

demo.queue(max_size=8, default_concurrency_limit=1)

if __name__ == "__main__":
    demo.launch()