Initial deploy: CrossKEY interactive 3D matching demo

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

README.md

@@ -1,12 +1,17 @@
 ---
 title: CrossKEY
-emoji: 📚
-colorFrom: yellow
-colorTo: pink
+emoji: "\U0001f9e0"
+colorFrom: blue
+colorTo: red
 sdk: gradio
-sdk_version: 6.12.0
+sdk_version: "5.29.0"
 app_file: app.py
-pinned: false
+license: mit
+short_description: 3D Cross-modal Keypoint Matching (MR-US)
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# CrossKEY Demo
+
+Interactive 3D visualization of cross-modal keypoint matching between MRI and intraoperative ultrasound.
+
+[Paper (IEEE TMI)](https://doi.org/10.1109/TMI.2026.3680352) | [Code (GitHub)](https://github.com/morozovdd/CrossKEY) | [arXiv](https://arxiv.org/abs/2507.18551)

app.py

@@ -0,0 +1,269 @@
+"""CrossKEY HuggingFace Space -- Interactive 3D Keypoint Matching Demo.
+
+Two-tab Gradio app:
+  Tab 1 (Explore): Pre-computed results with adjustable matching parameters.
+  Tab 2 (Your Data): Upload volumes + checkpoint for live inference.
+"""
+
+import logging
+import os
+import sys
+
+# Add space/ to path so local imports work both locally and on HF
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+import gradio as gr
+import numpy as np
+
+from inference import load_precomputed, run_inference, run_matching
+from visualization import build_matching_figure
+
+# ZeroGPU decorator -- no-op when running locally
+try:
+    import spaces
+    gpu_decorator = spaces.GPU
+except ImportError:
+    gpu_decorator = lambda fn: fn
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+logger = logging.getLogger("crosskey.app")
+
+# -- Load pre-computed data on startup --
+logger.info("Loading pre-computed demo data...")
+DEMO_DATA = load_precomputed(
+    os.path.join(os.path.dirname(os.path.abspath(__file__)), "precomputed")
+)
+logger.info(
+    "Loaded: %d MR descriptors, %d US descriptors",
+    len(DEMO_DATA["descriptors_mr"]),
+    len(DEMO_DATA["descriptors_us"]),
+)
+
+
+def update_demo(
+    ratio_threshold: float,
+    evaluation_threshold: float,
+    mutual: bool,
+    metric: str,
+) -> tuple:
+    """Re-run matching with new parameters and rebuild the figure."""
+    match_pairs, metrics = run_matching(
+        DEMO_DATA["descriptors_mr"],
+        DEMO_DATA["descriptors_us"],
+        DEMO_DATA["points_mr"],
+        DEMO_DATA["points_us"],
+        ratio_threshold=ratio_threshold,
+        mutual=mutual,
+        metric=metric,
+        evaluation_threshold=evaluation_threshold,
+    )
+
+    fig = build_matching_figure(
+        volume_mr=DEMO_DATA["volume_mr"],
+        volume_us=DEMO_DATA["volume_us"],
+        points_mr=DEMO_DATA["points_mr"],
+        points_us=DEMO_DATA["points_us"],
+        padded_shape_mr=tuple(DEMO_DATA["metadata"]["padded_shape_mr"]),
+        padded_shape_us=tuple(DEMO_DATA["metadata"]["padded_shape_us"]),
+        match_pairs=match_pairs,
+        metrics=metrics,
+        evaluation_threshold=evaluation_threshold,
+    )
+
+    return (
+        fig,
+        metrics['num_matches'],
+        metrics['num_correct'],
+        round(metrics['precision'], 1),
+        round(metrics['matching_score'], 4),
+    )
+
+
+@gpu_decorator
+def run_custom_inference(mr_file, us_file, heatmap_file, ckpt_file):
+    """Run inference on uploaded data. Uses ZeroGPU on HF Spaces."""
+    if any(f is None for f in [mr_file, us_file, heatmap_file, ckpt_file]):
+        raise gr.Error("Please upload all four files: MR volume, US volume, heatmap, and checkpoint.")
+
+    logger.info("Running inference on uploaded data...")
+    data = run_inference(
+        mr_path=mr_file.name,
+        us_path=us_file.name,
+        heatmap_path=heatmap_file.name,
+        checkpoint_path=ckpt_file.name,
+    )
+    return data
+
+
+def update_custom(
+    data: dict,
+    ratio_threshold: float,
+    evaluation_threshold: float,
+    mutual: bool,
+    metric: str,
+) -> tuple:
+    """Re-run matching on custom data with new parameters."""
+    if data is None:
+        raise gr.Error("Run inference first.")
+
+    match_pairs, metrics = run_matching(
+        data["descriptors_mr"],
+        data["descriptors_us"],
+        data["points_mr"],
+        data["points_us"],
+        ratio_threshold=ratio_threshold,
+        mutual=mutual,
+        metric=metric,
+        evaluation_threshold=evaluation_threshold,
+    )
+
+    fig = build_matching_figure(
+        volume_mr=data["volume_mr"],
+        volume_us=data["volume_us"],
+        points_mr=data["points_mr"],
+        points_us=data["points_us"],
+        padded_shape_mr=tuple(data["metadata"]["padded_shape_mr"]),
+        padded_shape_us=tuple(data["metadata"]["padded_shape_us"]),
+        match_pairs=match_pairs,
+        metrics=metrics,
+        evaluation_threshold=evaluation_threshold,
+    )
+
+    return (
+        fig,
+        metrics['num_matches'],
+        metrics['num_correct'],
+        round(metrics['precision'], 1),
+        round(metrics['matching_score'], 4),
+    )
+
+
+# -- Build Gradio UI --
+
+with gr.Blocks(
+    title="CrossKEY -- 3D Cross-modal Keypoint Matching",
+) as demo:
+    gr.Markdown(
+        "# CrossKEY\n"
+        "**3D Cross-modal Keypoint Descriptor for MR-US Matching and Registration**"
+    )
+
+    with gr.Tabs():
+        # ---- Tab 1: Explore ----
+        with gr.Tab("Explore"):
+            with gr.Row():
+                with gr.Column(scale=1, min_width=260):
+                    gr.Markdown("### Matching Parameters")
+                    demo_ratio = gr.Slider(
+                        minimum=0.5, maximum=1.0, value=0.75, step=0.05,
+                        label="Ratio Threshold",
+                    )
+                    demo_eval_thresh = gr.Slider(
+                        minimum=1.0, maximum=10.0, value=5.0, step=0.5,
+                        label="Evaluation Threshold (mm)",
+                    )
+                    demo_mutual = gr.Checkbox(value=True, label="Mutual Nearest Neighbor")
+                    demo_metric = gr.Dropdown(
+                        choices=["euclidean", "cosine"], value="euclidean",
+                        label="Distance Metric",
+                    )
+
+                    gr.Markdown("### Results")
+                    with gr.Group():
+                        with gr.Row():
+                            demo_n_matches = gr.Number(label="Matches", interactive=False)
+                            demo_n_correct = gr.Number(label="Correct", interactive=False)
+                        with gr.Row():
+                            demo_precision = gr.Number(label="Precision (%)", interactive=False)
+                            demo_match_score = gr.Number(label="Match Score", interactive=False)
+
+                with gr.Column(scale=3):
+                    demo_plot = gr.Plot(label="3D Matching Visualization")
+
+            demo_inputs = [demo_ratio, demo_eval_thresh, demo_mutual, demo_metric]
+            demo_outputs = [demo_plot, demo_n_matches, demo_n_correct, demo_precision, demo_match_score]
+
+            # Update on any parameter change
+            for inp in demo_inputs:
+                inp.change(fn=update_demo, inputs=demo_inputs, outputs=demo_outputs)
+
+            # Load initial results
+            demo.load(fn=update_demo, inputs=demo_inputs, outputs=demo_outputs)
+
+        # ---- Tab 2: Your Data ----
+        with gr.Tab("Your Data"):
+            gr.Markdown(
+                "Upload your own MR volume, US volume, heatmap, and a trained CrossKEY checkpoint.\n\n"
+                "Inference runs on GPU and may take 30-60 seconds."
+            )
+
+            with gr.Row():
+                custom_mr = gr.File(label="MR Volume (.nii.gz)", file_types=[".nii.gz"])
+                custom_us = gr.File(label="US Volume (.nii.gz)", file_types=[".nii.gz"])
+            with gr.Row():
+                custom_heatmap = gr.File(label="Heatmap (.nii.gz)", file_types=[".nii.gz"])
+                custom_ckpt = gr.File(label="Checkpoint (.ckpt)", file_types=[".ckpt"])
+
+            custom_run_btn = gr.Button("Run Inference (GPU)", variant="primary")
+
+            with gr.Row():
+                with gr.Column(scale=1, min_width=260):
+                    gr.Markdown("### Matching Parameters")
+                    custom_ratio = gr.Slider(
+                        minimum=0.5, maximum=1.0, value=0.75, step=0.05,
+                        label="Ratio Threshold",
+                    )
+                    custom_eval_thresh = gr.Slider(
+                        minimum=1.0, maximum=10.0, value=5.0, step=0.5,
+                        label="Evaluation Threshold (mm)",
+                    )
+                    custom_mutual = gr.Checkbox(value=True, label="Mutual Nearest Neighbor")
+                    custom_metric = gr.Dropdown(
+                        choices=["euclidean", "cosine"], value="euclidean",
+                        label="Distance Metric",
+                    )
+
+                    gr.Markdown("### Results")
+                    with gr.Group():
+                        with gr.Row():
+                            custom_n_matches = gr.Number(label="Matches", interactive=False)
+                            custom_n_correct = gr.Number(label="Correct", interactive=False)
+                        with gr.Row():
+                            custom_precision = gr.Number(label="Precision (%)", interactive=False)
+                            custom_match_score = gr.Number(label="Match Score", interactive=False)
+
+                with gr.Column(scale=3):
+                    custom_plot = gr.Plot(label="3D Matching Visualization")
+
+            # State to hold inference results
+            custom_data_state = gr.State(value=None)
+
+            custom_param_inputs = [custom_ratio, custom_eval_thresh, custom_mutual, custom_metric]
+            custom_outputs = [custom_plot, custom_n_matches, custom_n_correct, custom_precision, custom_match_score]
+
+            # Inference button: run model, then update visualization
+            def infer_and_display(mr_file, us_file, heatmap_file, ckpt_file, ratio, eval_thresh, mutual, metric):
+                data = run_custom_inference(mr_file, us_file, heatmap_file, ckpt_file)
+                fig, n_m, n_c, prec, ms = update_custom(data, ratio, eval_thresh, mutual, metric)
+                return data, fig, n_m, n_c, prec, ms
+
+            custom_run_btn.click(
+                fn=infer_and_display,
+                inputs=[custom_mr, custom_us, custom_heatmap, custom_ckpt] + custom_param_inputs,
+                outputs=[custom_data_state] + custom_outputs,
+            )
+
+            # Re-match on parameter change (no re-inference)
+            for inp in custom_param_inputs:
+                inp.change(
+                    fn=update_custom,
+                    inputs=[custom_data_state] + custom_param_inputs,
+                    outputs=custom_outputs,
+                )
+
+
+if __name__ == "__main__":
+    demo.launch(
+        theme=gr.themes.Soft(),
+        css="footer {display: none !important;}",
+    )

inference.py

@@ -0,0 +1,144 @@
+"""Descriptor extraction and matching for CrossKEY HF Space.
+
+Provides functions for:
+1. Re-running KNN matching with new parameters (CPU, fast)
+2. Full inference from uploaded volumes + checkpoint (GPU)
+"""
+
+import json
+from pathlib import Path
+from typing import Dict, List, Tuple
+
+import numpy as np
+import torch
+
+from src.data.datamodule import DescriptorDataModule
+from src.model.descriptor import Descriptor
+from src.model.matcher import KNNMatcher
+from src.utils.utils import load_nifti
+
+from visualization import downsample_volume
+
+
+def load_precomputed(precomputed_dir: str = "precomputed") -> dict:
+    """Load all pre-computed data for the default demo tab.
+
+    Returns:
+        Dict with keys: descriptors_mr, descriptors_us, points_mr, points_us,
+                        volume_mr, volume_us, metadata
+    """
+    d = Path(precomputed_dir)
+    with open(d / "metadata.json") as f:
+        metadata = json.load(f)
+
+    return {
+        "descriptors_mr": torch.load(d / "descriptors_mr.pt", weights_only=True),
+        "descriptors_us": torch.load(d / "descriptors_us.pt", weights_only=True),
+        "points_mr": torch.load(d / "points_mr.pt", weights_only=True).numpy(),
+        "points_us": torch.load(d / "points_us.pt", weights_only=True).numpy(),
+        "volume_mr": np.load(d / "volume_mr.npy"),
+        "volume_us": np.load(d / "volume_us.npy"),
+        "metadata": metadata,
+    }
+
+
+def run_matching(
+    descriptors_mr: torch.Tensor,
+    descriptors_us: torch.Tensor,
+    points_mr: np.ndarray,
+    points_us: np.ndarray,
+    ratio_threshold: float = 0.75,
+    mutual: bool = True,
+    metric: str = "euclidean",
+    evaluation_threshold: float = 5.0,
+) -> Tuple[List[Tuple[int, int, float]], Dict[str, float]]:
+    """Run KNN matching with given parameters. CPU-only, fast (<1s).
+
+    Returns:
+        (match_pairs, metrics) -- same format as KNNMatcher.match_and_evaluate()
+    """
+    matcher = KNNMatcher(
+        k=1,
+        distance_threshold=float("inf"),
+        ratio_threshold=ratio_threshold,
+        mutual=mutual,
+        metric=metric,
+        evaluation_threshold=evaluation_threshold,
+    )
+    return matcher.match_and_evaluate(
+        descriptors_mr, descriptors_us, points_mr, points_us,
+    )
+
+
+def run_inference(
+    mr_path: str,
+    us_path: str,
+    heatmap_path: str,
+    checkpoint_path: str,
+    batch_size: int = 64,
+    grid_spacing: int = 8,
+) -> dict:
+    """Run full inference on uploaded volumes. Requires GPU.
+
+    Args:
+        mr_path: Path to uploaded MR NIfTI file.
+        us_path: Path to uploaded US NIfTI file.
+        heatmap_path: Path to uploaded heatmap NIfTI file.
+        checkpoint_path: Path to uploaded checkpoint.
+        batch_size: Inference batch size.
+        grid_spacing: Grid spacing for US keypoint generation.
+
+    Returns:
+        Dict with same keys as load_precomputed().
+    """
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Load model
+    model = Descriptor.load_from_checkpoint(checkpoint_path)
+    model.eval()
+    model.to(device)
+
+    # Create datamodule with custom paths
+    dm = DescriptorDataModule(
+        data_dir=".",  # Not used when paths are specified
+        batch_size=batch_size,
+        num_workers=0,
+        patch_size=(32, 32, 32),
+        grid_spacing=grid_spacing,
+        mr_path=mr_path,
+        us_path=us_path,
+        heatmap_path=heatmap_path,
+    )
+    dm.setup(stage="test")
+
+    # Extract descriptors
+    all_desc, all_pts, all_mod = [], [], []
+    with torch.no_grad():
+        for batch in dm.test_dataloader():
+            desc = model(batch["patch"].to(device))
+            all_desc.append(desc.cpu())
+            all_pts.append(batch["point"].cpu())
+            all_mod.extend(batch["modality"])
+
+    all_desc = torch.cat(all_desc)
+    all_pts = torch.cat(all_pts)
+    mr_mask = torch.tensor([m == "mr" for m in all_mod])
+
+    # Downsample volumes for rendering
+    mr_vol = load_nifti(mr_path)
+    us_vol = load_nifti(us_path)
+    mr_norm = (mr_vol - mr_vol.min()) / (mr_vol.max() - mr_vol.min() + 1e-8)
+    us_norm = (us_vol - us_vol.min()) / (us_vol.max() - us_vol.min() + 1e-8)
+
+    return {
+        "descriptors_mr": all_desc[mr_mask],
+        "descriptors_us": all_desc[~mr_mask],
+        "points_mr": all_pts[mr_mask].numpy(),
+        "points_us": all_pts[~mr_mask].numpy(),
+        "volume_mr": downsample_volume(mr_norm),
+        "volume_us": downsample_volume(us_norm),
+        "metadata": {
+            "padded_shape_mr": list(dm._mr_volume.shape),
+            "padded_shape_us": list(dm._us_volume.shape),
+        },
+    }

precomputed/descriptors_mr.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9507f55c6c0666346b8ab93a954c89e8df192d937fbc18452a2eacc5b47e1fd2
+size 2098778

precomputed/descriptors_us.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:393b7b33a75392322615d4523c1b4b6ff7b6664491b2d060aef9df95e317952e
+size 4443738

precomputed/metadata.json

@@ -0,0 +1,16 @@
+{
+  "padded_shape_mr": [
+    224,
+    224,
+    198
+  ],
+  "padded_shape_us": [
+    224,
+    224,
+    198
+  ],
+  "num_mr_descriptors": 1024,
+  "num_us_descriptors": 2169,
+  "descriptor_dim": 512,
+  "grid_spacing": 8
+}

precomputed/points_mr.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:86dee258740fe89673a551a7a1611f39b9566c5e6071521f7424ff0a71ecec54
+size 13879

precomputed/points_us.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a6b8d0ece568a4552e6f7f88c37d688ffeb2daf8ab559baaab128e826d28e0f2
+size 27575

precomputed/volume_mr.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:feb4d465b3e378693f731a27e0d53544db16ddac9e5e507b5ca073f6531e1bf0
+size 1048704

precomputed/volume_us.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:42fb079ee925a7a71d8c76a41c5a69963a67c2b7e585e747d7414c4a17f0162c
+size 1048704

requirements.txt

@@ -0,0 +1,4 @@
+crosskey @ git+https://github.com/morozovdd/CrossKEY.git
+gradio>=5.0,<6.0
+plotly>=6.0,<7.0
+scikit-image>=0.24,<1.0

visualization.py

@@ -0,0 +1,223 @@
+"""Plotly 3D visualization for CrossKEY matching results.
+
+Builds side-by-side volume isosurfaces with keypoints and match lines.
+MR volume on the left, US volume on the right, offset along the X axis.
+"""
+
+import numpy as np
+import plotly.graph_objects as go
+from scipy.ndimage import zoom
+from skimage.measure import marching_cubes
+
+
+def downsample_volume(volume: np.ndarray, target_size: int = 64) -> np.ndarray:
+    """Downsample volume to target_size^3 for browser-friendly rendering."""
+    factors = [target_size / s for s in volume.shape]
+    return zoom(volume, factors, order=1).astype(np.float32)
+
+
+def scale_points(
+    points: np.ndarray,
+    padded_shape: tuple,
+    volume_shape: tuple,
+) -> np.ndarray:
+    """Scale point coordinates from padded volume space to downsampled volume space."""
+    scale = np.array(volume_shape, dtype=float) / np.array(padded_shape, dtype=float)
+    return points * scale
+
+
+def create_isosurface_trace(
+    volume: np.ndarray,
+    level: float,
+    colorscale: str = "Gray",
+    opacity: float = 0.15,
+    name: str = "",
+    offset_x: float = 0.0,
+) -> go.Mesh3d:
+    """Create a Mesh3d trace from a volume isosurface via marching cubes.
+
+    Uses vertex intensity from the original volume for natural coloring.
+    """
+    verts, faces, _, _ = marching_cubes(volume, level=level)
+    # Sample volume intensity at each vertex for natural coloring
+    vi = np.clip(verts.astype(int), 0, np.array(volume.shape) - 1)
+    intensities = volume[vi[:, 0], vi[:, 1], vi[:, 2]]
+    # Axis remap: data (0,1,2) -> Plotly (z,x,y); negate axis 0 so cone points up
+    return go.Mesh3d(
+        x=verts[:, 1] + offset_x,
+        y=verts[:, 2],
+        z=-verts[:, 0],
+        i=faces[:, 0],
+        j=faces[:, 1],
+        k=faces[:, 2],
+        intensity=intensities,
+        colorscale=colorscale,
+        opacity=opacity,
+        name=name,
+        showlegend=True,
+        showscale=False,
+    )
+
+
+def create_keypoint_trace(
+    points: np.ndarray,
+    color: str,
+    size: float = 3.0,
+    opacity: float = 1.0,
+    name: str = "",
+    offset_x: float = 0.0,
+) -> go.Scatter3d:
+    """Create Scatter3d markers for keypoints."""
+    # Axis remap: data (0,1,2) -> Plotly (z,x,y); negate axis 0
+    return go.Scatter3d(
+        x=points[:, 1] + offset_x,
+        y=points[:, 2],
+        z=-points[:, 0],
+        mode="markers",
+        marker=dict(size=size, color=color, opacity=opacity),
+        name=name,
+        showlegend=True,
+    )
+
+
+def create_match_lines(
+    src_pts: np.ndarray,
+    tgt_pts: np.ndarray,
+    color: str,
+    width: float = 2.0,
+    name: str = "",
+    offset_x: float = 0.0,
+) -> go.Scatter3d:
+    """Create lines connecting matched source points to offset target points."""
+    # Axis remap: data (0,1,2) -> Plotly (z,x,y); negate axis 0
+    lx, ly, lz = [], [], []
+    for s, t in zip(src_pts, tgt_pts):
+        lx.extend([float(s[1]), float(t[1]) + offset_x, None])
+        ly.extend([float(s[2]), float(t[2]), None])
+        lz.extend([-float(s[0]), -float(t[0]), None])
+    return go.Scatter3d(
+        x=lx, y=ly, z=lz,
+        mode="lines",
+        line=dict(color=color, width=width),
+        name=name,
+        showlegend=True,
+    )
+
+
+def build_matching_figure(
+    volume_mr: np.ndarray,
+    volume_us: np.ndarray,
+    points_mr: np.ndarray,
+    points_us: np.ndarray,
+    padded_shape_mr: tuple,
+    padded_shape_us: tuple,
+    match_pairs: list,
+    metrics: dict,
+    evaluation_threshold: float = 5.0,
+    mr_level: float = 0.3,
+    us_level: float = 0.1,
+) -> go.Figure:
+    """Build the full 3D matching visualization."""
+    fig = go.Figure()
+
+    # Scale keypoints to match downsampled volume coordinates
+    pts_mr_viz = scale_points(points_mr, padded_shape_mr, volume_mr.shape)
+    pts_us_viz = scale_points(points_us, padded_shape_us, volume_us.shape)
+
+    # Side-by-side offset along Plotly x (= data axis 1)
+    gap = volume_mr.shape[1] * 0.3
+    offset_x = volume_mr.shape[1] + gap
+
+    # Volume isosurfaces with natural intensity coloring
+    try:
+        fig.add_trace(create_isosurface_trace(
+            volume_mr, level=mr_level, colorscale="Gray",
+            opacity=0.15, name="MR Surface",
+        ))
+    except ValueError:
+        pass
+
+    try:
+        fig.add_trace(create_isosurface_trace(
+            volume_us, level=us_level, colorscale="Hot",
+            opacity=0.15, name="US Surface", offset_x=offset_x,
+        ))
+    except ValueError:
+        pass
+
+    # Process matches
+    src_indices = [p[0] for p in match_pairs]
+    tgt_indices = [p[1] for p in match_pairs]
+
+    if match_pairs:
+        mr_matched = points_mr[src_indices]
+        us_matched = points_us[tgt_indices]
+        spatial_dist = np.linalg.norm(mr_matched - us_matched, axis=1)
+        correct = spatial_dist < evaluation_threshold
+
+        mr_matched_viz = pts_mr_viz[src_indices]
+        us_matched_viz = pts_us_viz[tgt_indices]
+
+        if correct.any():
+            fig.add_trace(create_match_lines(
+                mr_matched_viz[correct], us_matched_viz[correct],
+                color="rgba(0,200,0,0.6)", width=2,
+                name=f"Correct ({correct.sum()})", offset_x=offset_x,
+            ))
+
+        if (~correct).any():
+            fig.add_trace(create_match_lines(
+                mr_matched_viz[~correct], us_matched_viz[~correct],
+                color="rgba(255,0,0,0.3)", width=1,
+                name=f"Incorrect ({(~correct).sum()})", offset_x=offset_x,
+            ))
+
+        fig.add_trace(create_keypoint_trace(
+            mr_matched_viz, color="royalblue", size=4,
+            name=f"MR Matched ({len(mr_matched_viz)})",
+        ))
+        fig.add_trace(create_keypoint_trace(
+            us_matched_viz, color="crimson", size=4,
+            name=f"US Matched ({len(us_matched_viz)})", offset_x=offset_x,
+        ))
+
+    # Unmatched keypoints (faded)
+    matched_mr_set = set(src_indices)
+    matched_us_set = set(tgt_indices)
+    unmatched_mr = np.array([i not in matched_mr_set for i in range(len(pts_mr_viz))])
+    unmatched_us = np.array([i not in matched_us_set for i in range(len(pts_us_viz))])
+
+    if unmatched_mr.any():
+        fig.add_trace(create_keypoint_trace(
+            pts_mr_viz[unmatched_mr], color="royalblue",
+            size=1.5, opacity=0.2, name="MR Unmatched",
+        ))
+    if unmatched_us.any():
+        fig.add_trace(create_keypoint_trace(
+            pts_us_viz[unmatched_us], color="crimson",
+            size=1.5, opacity=0.2, name="US Unmatched", offset_x=offset_x,
+        ))
+
+    # Layout -- no fixed width so Plotly fills the Gradio container
+    fig.update_layout(
+        scene=dict(
+            xaxis=dict(visible=False),
+            yaxis=dict(visible=False),
+            zaxis=dict(visible=False),
+            aspectmode="data",
+            camera=dict(
+                up=dict(x=0, y=0, z=1),
+                eye=dict(x=0, y=-1.8, z=0.3),
+            ),
+        ),
+        height=700,
+        margin=dict(l=0, r=0, t=40, b=0),
+        legend=dict(
+            yanchor="top", y=0.99,
+            xanchor="left", x=0.01,
+            bgcolor="rgba(0,0,0,0.5)",
+            font=dict(color="white"),
+        ),
+    )
+
+    return fig