Spaces:
Running on Zero
Running on Zero
| """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 | |