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

try:
    import spaces
except ImportError:
    class _LocalSpaces:
        """Use the same app locally when the ZeroGPU helper is unavailable."""

        @staticmethod
        def GPU(*_args, **_kwargs):
            return lambda function: function

    spaces = _LocalSpaces()

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"))

PRECOMPUTED_ROOT = PROJECT_ROOT / "examples" / "precomputed"
PRECOMPUTED_EXAMPLES = {
    "0000c32f": {
        "label": "HY3D Example 1 · 0000c32f",
        "vertices": 138_148,
        "faces": 276_596,
    },
    "007e1061": {
        "label": "HY3D Example 2 · 007e1061",
        "vertices": 229_427,
        "faces": 459_490,
    },
    "0005df57": {
        "label": "HY3D Example 3 · 0005df57",
        "vertices": 275_363,
        "faces": 551_188,
    },
}
DEFAULT_EXAMPLE_ID = next(iter(PRECOMPUTED_EXAMPLES))


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 load_precomputed_example(example_id: str):
    """Return browser previews and original meshes for a bundled result."""
    if example_id not in PRECOMPUTED_EXAMPLES:
        raise gr.Error("Select one of the bundled examples.")

    example = PRECOMPUTED_EXAMPLES[example_id]
    example_dir = PRECOMPUTED_ROOT / example_id
    input_glb = example_dir / "input.glb"
    output_glb = example_dir / "mesh_mapped_v8.glb"
    input_ply = example_dir / "input.ply"
    output_ply = example_dir / "mesh_mapped_v8.ply"
    required = (input_glb, output_glb, input_ply, output_ply)
    if not all(path.is_file() for path in required):
        raise gr.Error(f"Bundled files are incomplete for {example['label']}.")

    summary = (
        f"**{example['label']}** — "
        f"{example['vertices']:,} vertices · {example['faces']:,} faces. "
        "The right viewer shows the final mesh segmented result."
    )
    return (
        summary,
        str(input_glb),
        str(output_glb),
        str(input_ply),
        str(output_ply),
    )


@spaces.GPU(duration=300)
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


default_example = load_precomputed_example(DEFAULT_EXAMPLE_ID)

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.
"""
    )

    gr.Markdown(
        """
## Precomputed examples

Explore three bundled HY3D inputs and their final HiTOPS mesh-mapping results
immediately. Switching examples does not run the pipeline or consume ZeroGPU quota.
"""
    )

    example_selector = gr.Radio(
        choices=[
            (example["label"], example_id)
            for example_id, example in PRECOMPUTED_EXAMPLES.items()
        ],
        value=DEFAULT_EXAMPLE_ID,
        label="Select a precomputed example",
    )
    example_summary = gr.Markdown(default_example[0])

    with gr.Row():
        example_input_preview = gr.Model3D(
            value=default_example[1],
            label="Input mesh",
            height=440,
            interactive=False,
            clear_color=(0.08, 0.08, 0.10, 1.0),
        )
        example_output_preview = gr.Model3D(
            value=default_example[2],
            label="HiTOPS mesh mapping result",
            height=440,
            interactive=False,
            clear_color=(0.08, 0.08, 0.10, 1.0),
        )

    with gr.Row():
        example_input_download = gr.File(
            value=default_example[3],
            label="Download input PLY",
            interactive=False,
        )
        example_output_download = gr.File(
            value=default_example[4],
            label="Download mapped result PLY",
            interactive=False,
        )

    example_selector.change(
        fn=load_precomputed_example,
        inputs=example_selector,
        outputs=[
            example_summary,
            example_input_preview,
            example_output_preview,
            example_input_download,
            example_output_download,
        ],
        concurrency_limit=4,
        api_name=False,
    )

    gr.Markdown(
        """
---

## Run HiTOPS on your mesh
"""
    )

    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",
    )

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

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