engine.py

"""MEOW sB25 inference and per-request exports; no web server or global history."""
from __future__ import annotations

import json
import os
from pathlib import Path
import shutil
import sys
import tempfile
import time

import numpy as np
from PIL import Image, ImageOps

ROOT = Path(__file__).resolve().parent
VENDOR = ROOT / "vendor"
for path in (VENDOR, VENDOR / "third_party/uniception", VENDOR / "scripts"):
    sys.path.insert(0, str(path))

MAX_IMAGES = 8
MAX_PIXELS = 25_000_000
MODEL_REPO = "qgli/MEOW-4D-Space-Reconstruction"
MODEL_FILE = "MEOW-sB25.pth"


def result_directory():
    """Expire original export files as well as Gradio's own copied-file cache."""
    root = Path(tempfile.gettempdir()) / "meow-space-results"
    root.mkdir(exist_ok=True)
    cutoff = time.time() - 3600
    for child in root.iterdir():
        if child.is_dir() and not child.is_symlink() and child.stat().st_mtime < cutoff:
            shutil.rmtree(child)
    return Path(tempfile.mkdtemp(prefix="result-", dir=root))


def prepare_images(files, workdir):
    """Normalise orientation once; avoid metadata and user names in exports."""
    from erp_detect import detect_full_erp

    if not files or not 2 <= len(files) <= MAX_IMAGES:
        raise ValueError("Upload 2–8 overlapping images of the same scene.")
    workdir = Path(workdir)
    workdir.mkdir(parents=True, exist_ok=True)
    paths, metadata = [], []
    for i, source in enumerate(files):
        with Image.open(source) as original:
            if original.width * original.height > MAX_PIXELS:
                raise ValueError("Each image must contain at most 25 megapixels.")
            image = ImageOps.exif_transpose(original).convert("RGB")
            width, height = image.size
            is_pano, _ = detect_full_erp(np.asarray(image))
            target = workdir / f"view_{i:02d}.png"
            # This file is consumed immediately, not downloaded. PNG compression
            # spends seconds recompressing multi-megapixel uploads for no gain;
            # level 0 is lossless and leaves the model's pixels unchanged.
            image.save(target, compress_level=0)
        paths.append(str(target))
        metadata.append({"view": i, "width": width, "height": height,
                         "detected_full_panorama": bool(is_pano),
                         "source_aspect": 2.0 if is_pano else width / height})
    return paths, metadata


def load_model(device="cuda"):
    import torch
    from huggingface_hub import hf_hub_download
    from mapanything.models import _init_hydra_config, init_model

    checkpoint = os.environ.get("MEOW_CHECKPOINT")
    if not checkpoint:
        checkpoint = hf_hub_download(MODEL_REPO, MODEL_FILE,
                                     token=os.environ.get("HF_TOKEN"))
    config = _init_hydra_config("configs/train.yaml", overrides=[
        "model=mapanything", "machine=default", "model/task=aug_training",
        "++model.task.ar_prob=1.0", "++model.model_config.pano_wrap_dpt=true",
    ])
    # Allocate shapes only, then assign the complete released FP32 state.
    # This avoids holding a second 5 GB randomly initialised model in CPU RAM.
    with torch.device("meta"):
        model = init_model(config.model.model_str, config.model.model_config)
    state = torch.load(checkpoint, map_location="cpu", weights_only=True, mmap=True)
    model.load_state_dict(state["model"], strict=True, assign=True)
    del state
    return model.eval().to(device)


def camera_segments(poses, scale):
    """Symbolic frusta, not calibrated camera models (ERP has no pinhole frustum)."""
    result = []
    corners = np.array([[-.6, -.4, 1], [.6, -.4, 1],
                        [.6, .4, 1], [-.6, .4, 1]]) * scale
    for pose in poses:
        world = corners @ pose[:3, :3].T + pose[:3, 3]
        origin = pose[:3, 3]
        result.extend([[origin, corner] for corner in world])
        result.extend([[world[i], world[(i + 1) % 4]] for i in range(4)])
    return np.asarray(result)


def export_scene(points, colors, poses, outdir, footprints=None):
    import trimesh
    from camera_glyphs import COLORS, wire_mesh

    outdir = Path(outdir)
    # PLY and cameras.json retain the original metric/OpenCV world coordinates.
    cloud = trimesh.points.PointCloud(points, colors=colors)
    ply = outdir / "pointcloud.ply"
    cloud.export(ply)
    scene = trimesh.Scene()
    scene.add_geometry(cloud, node_name="pointcloud")
    extent = np.percentile(points, 95, axis=0) - np.percentile(points, 5, axis=0)
    scale = max(float(np.linalg.norm(extent)) * .025, .01)
    # Keep the compact size of the new demo; never use tan(FOV/2) to enlarge it.
    for index, pose in enumerate(poses):
        if footprints is not None:
            segments, info = footprints[index]
            segments = segments * scale
            segments = segments @ pose[:3, :3].T + pose[:3, 3]
            color = COLORS[info["family_hint"]]
        else:
            segments = camera_segments([pose], scale)
            color = [255, 180, 60, 255]
        wire = wire_mesh(segments, scale * .009, color)
        if wire is not None:
            scene.add_geometry(wire, node_name=f"camera_{index}")
        center = trimesh.creation.icosphere(subdivisions=1, radius=scale * .06)
        center.apply_translation(pose[:3, 3])
        center.visual.vertex_colors = color
        scene.add_geometry(center, node_name=f"camera_center_{index}")
    # Viewer only: rotate OpenCV Y-down to Y-up; export the transform explicitly.
    display_transform = np.diag([1., -1., -1., 1.])
    scene.apply_transform(display_transform)
    glb = outdir / "scene.glb"
    scene.export(glb)
    return str(glb), str(ply), display_transform.tolist()


def reconstruct(model, files, confidence_percentile=10, max_points=180_000):
    import torch
    from mapanything.utils.image import load_images, rgb
    from camera_glyphs import ray_footprint

    started = time.perf_counter()
    outdir = result_directory()
    # Inputs live in a separate temporary directory which is removed immediately.
    with tempfile.TemporaryDirectory(prefix="meow-input-") as inputs:
        paths, metadata = prepare_images(files, inputs)
        views = load_images(paths, verbose=False, resize_mode="fixed_squeeze")
        if len(views) != len(paths):
            raise ValueError("An input image could not be decoded.")
        image_colors = [np.uint8(np.clip(rgb(v["img"][0], "dinov2"), 0, 1) * 255) for v in views]
        for view, meta in zip(views, metadata):
            view["aspect_ratio"] = torch.tensor([meta["source_aspect"]], dtype=torch.float32)
            view["pano_wrap"] = meta["detected_full_panorama"]
            for key, value in view.items():
                if torch.is_tensor(value):
                    view[key] = value.to("cuda")
        torch.cuda.synchronize()
        prepared = time.perf_counter()
        with torch.inference_mode():
            predictions = model.infer(views, memory_efficient_inference=False,
                                      use_amp=True, amp_dtype="bf16",
                                      apply_mask=True, mask_edges=True)
        torch.cuda.synchronize()
        inferred = time.perf_counter()
        points, colors, poses, footprints = [], [], [], []
        for prediction, color, meta in zip(predictions, image_colors, metadata):
            xyz = prediction["pts3d"][0].float().cpu().numpy().reshape(-1, 3)
            conf = prediction["conf"][0].float().cpu().numpy().reshape(-1)
            valid = np.isfinite(xyz).all(axis=1) & np.isfinite(conf) & (np.linalg.norm(xyz, axis=1) > 0)
            if valid.any():
                valid &= conf >= np.percentile(conf[valid], confidence_percentile)
            points.append(xyz[valid])
            colors.append(color.reshape(-1, 3)[valid])
            poses.append(prediction["camera_poses"][0].float().cpu().numpy())
            rays = prediction["ray_directions"][0].float().cpu().numpy()
            mask = prediction.get("non_ambiguous_mask")
            mask = mask[0].cpu().numpy() if mask is not None else None
            footprint = ray_footprint(rays, mask, 1.)
            footprints.append(footprint)
            meta.update(footprint[1])
        del predictions, views
    points, colors, poses = np.concatenate(points), np.concatenate(colors), np.stack(poses)
    if not len(points):
        raise ValueError("No valid points survived filtering; lower the confidence filter.")
    # Deterministic display/export budget; changes no model predictions.
    if len(points) > max_points:
        selected = np.linspace(0, len(points) - 1, max_points, dtype=int)
        points, colors = points[selected], colors[selected]
    glb, ply, transform = export_scene(points, colors, poses, outdir, footprints)
    camera_file = outdir / "cameras.json"
    camera_file.write_text(json.dumps({
        "model": "MEOW sB25", "units": "metres",
        "pose_convention": "camera-to-world, OpenCV (+X right, +Y down, +Z forward)",
        "glb_display_from_model": transform,
        "camera_glyphs": "wire footprints from predicted rays; compact symbolic radius; family hints are heuristic",
        "views": [dict(meta, camera_to_world=pose.tolist()) for meta, pose in zip(metadata, poses)],
    }, indent=2), encoding="utf-8")
    elapsed = time.perf_counter() - started
    summary = (f"{len(files)} views · {len(points):,} displayed/exported points · {elapsed:.1f} s\n\n"
               f"Image preparation {prepared - started:.2f} s · "
               f"GPU reconstruction {inferred - prepared:.2f} s · "
               f"filtering/export {elapsed - (inferred - started):.2f} s. "
               "Upload, queueing and browser download/rendering are additional.\n\n"
               "Full field of view retained. Panorama flags are inferred from pixels. "
               "Camera wires follow predicted rays: yellow perspective-like, teal fisheye-like, purple ERP-like. "
               "Family hints and angular extents are approximate; glyph size is symbolic. "
               "Confidence filtering affects the visualisation only; this is not a benchmark score.\n\n" +
               " · ".join(f"View {m['view']}: {m['family_hint']} ({m['angular_extent_deg']}° extent)" for m in metadata))
    return glb, ply, str(camera_file), summary