Spaces:
Running on Zero
Running on Zero
| import spaces # MUST come before torch / any CUDA-touching import | |
| import os | |
| import sys | |
| import math | |
| import time | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| import gradio as gr | |
| from huggingface_hub import hf_hub_download | |
| # Make the bundled src/ package importable (SAMatcher modeling code). | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| import src # noqa: F401 (SAMatcher modeling package) | |
| from hydra.utils import instantiate | |
| from omegaconf import OmegaConf | |
| MODEL_REPO = "SSSSphinx/SAMatcher" | |
| CONFIG_PATH = os.path.join( | |
| os.path.dirname(os.path.abspath(__file__)), "src", "configs", | |
| "sam2.1", "sam2.1_hiera_l_sammatcher.yaml") | |
| INPUT_SIZE = 1024 | |
| STRIDE_ALIGN = 16 # eval scripts use align='disk' -> stride 16 on the long side | |
| def load_model(): | |
| """Instantiate SAM2HQBase from the authors' config and load the released checkpoint. | |
| The config is loaded directly with OmegaConf (hydra's pkg:// search path is | |
| environment-sensitive inside the Space container), then instantiated with | |
| hydra's recursive instantiate helper, exactly as the authors' build script does. | |
| """ | |
| cfg = OmegaConf.load(CONFIG_PATH) | |
| model = instantiate(cfg.model, _recursive_=True) | |
| model.eval() | |
| ckpt_path = hf_hub_download(MODEL_REPO, "weights/samatcher_best.ckpt") | |
| state = torch.load(ckpt_path, map_location="cpu", weights_only=False) | |
| missing, unexpected = model.load_state_dict(state, strict=False) | |
| # The 6 fuser logit_scale params exist only on the cosine-attention fallback | |
| # path; they are absent from the checkpoint by design, so ignore them. | |
| real_missing = [k for k in missing if not k.endswith("attn.logit_scale")] | |
| if real_missing or unexpected: | |
| raise RuntimeError( | |
| f"Checkpoint mismatch — missing={real_missing[:5]} unexpected={list(unexpected)[:5]}" | |
| ) | |
| return model | |
| print("Loading SAMatcher (this downloads ~960 MB of weights on first boot)...") | |
| model = load_model() | |
| model.to("cuda") | |
| print("SAMatcher loaded.") | |
| # --------------------------------------------------------------------------- | |
| # Preprocessing — mirrors dloc read_overlap_image(align='disk', resize=1024, | |
| # pad_square=True, RGB): long side -> 1024, stride-aligned to 16, zero-padded | |
| # to a 1024x1024 square at bottom/right, /255, NCHW. | |
| # --------------------------------------------------------------------------- | |
| def preprocess(image_rgb: np.ndarray): | |
| """Returns (model_input_tensor[1,3,S,S], content_h, content_w, scale_x, scale_y).""" | |
| h, w = image_rgb.shape[:2] | |
| scale = INPUT_SIZE / max(h, w) | |
| w_r, h_r = max(1, int(round(w * scale))), max(1, int(round(h * scale))) | |
| w_new = int(math.ceil(w_r / STRIDE_ALIGN) * STRIDE_ALIGN) | |
| h_new = int(math.ceil(h_r / STRIDE_ALIGN) * STRIDE_ALIGN) | |
| img = cv2.resize(image_rgb.astype(np.float32), (w_new, h_new), interpolation=cv2.INTER_LINEAR) | |
| padded = np.zeros((INPUT_SIZE, INPUT_SIZE, 3), dtype=np.float32) if max(w_new, h_new) <= INPUT_SIZE else None | |
| if padded is None: # extremely wide/tall images: hard-clip to 1024 | |
| padded = np.zeros((INPUT_SIZE, INPUT_SIZE, 3), dtype=np.float32) | |
| w_new, h_new = min(w_new, INPUT_SIZE), min(h_new, INPUT_SIZE) | |
| padded[:h_new, :w_new, :] = img[:h_new, :w_new, :] | |
| inp = torch.from_numpy(padded.transpose(2, 0, 1)[None] / 255.0).float() | |
| sx = float(w) / float(w_new) | |
| sy = float(h) / float(h_new) | |
| return inp, h_new, w_new, sx, sy | |
| # --------------------------------------------------------------------------- | |
| # Visualization — replicates the authors' evaluation plotting: | |
| # magenta soft-mask overlay (alpha 0.35, keeps brightness outside the mask) + | |
| # magenta contour + red bbox, and a cropped matching-window pair. | |
| # --------------------------------------------------------------------------- | |
| def render_overlay(image_rgb: np.ndarray, soft_mask: np.ndarray, box_xyxy) -> np.ndarray: | |
| img = image_rgb.astype(np.float32) / 255.0 | |
| m = np.clip(soft_mask, 0.0, 1.0).astype(np.float32) | |
| if m.shape != image_rgb.shape[:2]: | |
| m = cv2.resize(m, (image_rgb.shape[1], image_rgb.shape[0]), interpolation=cv2.INTER_LINEAR) | |
| out = img.copy() | |
| a = (0.35 * m)[..., None] | |
| magenta = np.array([1.0, 0.0, 1.0], dtype=np.float32) | |
| out = (1.0 - a) * out + a * magenta | |
| out_u8 = np.clip(out * 255.0, 0, 255).astype(np.uint8) | |
| m_bin = (m > 0.5).astype(np.uint8) | |
| if m_bin.sum() > 0: | |
| contours, _ = cv2.findContours(m_bin, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE) | |
| cv2.drawContours(out_u8, contours, -1, (255, 0, 255), 2, lineType=cv2.LINE_AA) | |
| if box_xyxy is not None: | |
| x0, y0, x1, y1 = [int(round(float(v))) for v in box_xyxy] | |
| h_img, w_img = out_u8.shape[:2] | |
| x0 = max(0, min(w_img - 1, x0)); x1 = max(0, min(w_img, x1)) | |
| y0 = max(0, min(h_img - 1, y0)); y1 = max(0, min(h_img, y1)) | |
| if x1 > x0 and y1 > y0: | |
| cv2.rectangle(out_u8, (x0, y0), (x1, y1), (255, 0, 0), max(2, w_img // 300)) | |
| return out_u8 | |
| def render_crop_pair(image_rgb: np.ndarray, box_xyxy, min_side: int = 224) -> np.ndarray: | |
| """Crop the predicted co-visible window; returns RGB crop (uint8).""" | |
| h, w = image_rgb.shape[:2] | |
| if box_xyxy is None: | |
| return image_rgb.copy() | |
| x0, y0, x1, y1 = [int(round(float(v))) for v in box_xyxy] | |
| x0 = max(0, min(w - 1, x0)); x1 = max(0, min(w, x1)) | |
| y0 = max(0, min(h - 1, y0)); y1 = max(0, min(h, y1)) | |
| if (x1 - x0) < 8 or (y1 - y0) < 8: | |
| return image_rgb.copy() | |
| crop = image_rgb[y0:y1, x0:x1] | |
| ch, cw = crop.shape[:2] | |
| if min(ch, cw) < min_side: # upscale tiny crops for visibility | |
| s = min_side / min(ch, cw) | |
| crop = cv2.resize(crop, (max(1, int(round(cw * s))), max(1, int(round(ch * s)))), | |
| interpolation=cv2.INTER_CUBIC) | |
| return np.clip(crop, 0, 255).astype(np.uint8) | |
| def run_pair(image0, image1, mask_thresh=0.3, use_boxes=True): | |
| """Full SAMatcher inference on one image pair. Returns numpy visualizations on CPU.""" | |
| imgs = [] | |
| metas = [] | |
| for image in (image0, image1): | |
| rgb = image if image.ndim == 3 else cv2.cvtColor(image, cv2.COLOR_GRAY2RGB) | |
| if rgb.shape[2] == 4: | |
| rgb = cv2.cvtColor(rgb, cv2.COLOR_RGBA2RGB) | |
| inp, ch, cw, sx, sy = preprocess(rgb) | |
| imgs.append(inp) | |
| metas.append((rgb, ch, cw, sx, sy)) | |
| batch = torch.cat([imgs[0], imgs[1]], dim=0).to("cuda") | |
| try: | |
| with torch.no_grad(): | |
| masks, (bbox0, bbox1) = model.forward(batch[0:1], batch[1:2]) | |
| finally: | |
| del batch, imgs | |
| return finalize(metas, masks, bbox0, bbox1, float(mask_thresh)) | |
| def finalize(metas, masks, bbox0, bbox1, mask_thresh): | |
| """Map raw outputs (padded-1024 coords) back to original image coords.""" | |
| outs = [] | |
| for idx, (rgb, ch, cw, sx, sy) in enumerate(metas): | |
| raw = masks[0, idx] # (256, 256) logits | |
| # Reference order: 256 -> padded 1024 -> crop content area -> original size | |
| raw = F.interpolate(raw[None, None], size=(INPUT_SIZE, INPUT_SIZE), | |
| mode="bilinear", align_corners=False) | |
| raw = raw[:, :, :ch, :cw] | |
| raw = F.interpolate(raw, size=rgb.shape[:2], mode="bilinear", | |
| align_corners=False) | |
| soft = raw[0, 0].sigmoid().cpu().numpy() | |
| # bbox: clamp to the un-padded content area, then scale to original coords | |
| box = bbox0[0].detach().cpu().float() if idx == 0 else bbox1[0].detach().cpu().float() | |
| box[0::2] = box[0::2].clamp(0, cw - 1) | |
| box[1::2] = box[1::2].clamp(0, ch - 1) | |
| box = (box * torch.tensor([sx, sy, sx, sy])).tolist() | |
| box[0::2] = [min(box[0], rgb.shape[1] - 1), min(box[2], rgb.shape[1] - 1)] | |
| box[1::2] = [min(box[1], rgb.shape[0] - 1), min(box[3], rgb.shape[0] - 1)] | |
| outs.append((rgb, soft, box)) | |
| return outs | |
| def match(image0, image1, mask_thresh=0.3, show_boxes=True, progress=gr.Progress(track_tqdm=True)): | |
| """Estimate the co-visible (matching) region of an image pair with SAMatcher. | |
| Args: | |
| image0: Reference view (RGB image). | |
| image1: Second view (RGB image; often a different zoom / scale). | |
| mask_thresh: Sigmoid threshold applied to the co-visibility mask (0.3 matches the | |
| paper's evaluation default; 0.4 is the value the tuned scripts use). | |
| show_boxes: Draw the predicted co-visible bounding boxes on both views. | |
| Returns: | |
| (overlay_view0, overlay_view1, crop0, crop1, info_text) | |
| """ | |
| t0 = time.time() | |
| if image0 is None or image1 is None: | |
| raise gr.Error("Please provide both images of the pair.") | |
| if isinstance(image0, str): | |
| image0 = cv2.cvtColor(cv2.imread(image0, cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB) | |
| if isinstance(image1, str): | |
| image1 = cv2.cvtColor(cv2.imread(image1, cv2.IMREAD_COLOR), cv2.COLOR_BGR2RGB) | |
| (rgb0, soft0, box0), (rgb1, soft1, box1) = run_pair(image0, image1, mask_thresh) | |
| ov0 = render_overlay(rgb0, soft0, box0 if show_boxes else None) | |
| ov1 = render_overlay(rgb1, soft1, box1 if show_boxes else None) | |
| crop0 = render_crop_pair(rgb0, box0) | |
| crop1 = render_crop_pair(rgb1, box1) | |
| # Co-visibility summary (per-view fractions; the two masks live in different | |
| # coordinate frames, so a cross-view IoU is not meaningful) | |
| a0 = (soft0 > mask_thresh) | |
| a1 = (soft1 > mask_thresh) | |
| ratio0 = float(a0.mean()) | |
| ratio1 = float(a1.mean()) | |
| dt = time.time() - t0 | |
| info = (f"Co-visible fraction — view A: {ratio0 * 100:.1f}% · view B: {ratio1 * 100:.1f}% · " | |
| f"inference: {dt:.2f}s") | |
| return ov0, ov1, crop0, crop1, info | |
| # --------------------------------------------------------------------------- | |
| # UI | |
| # --------------------------------------------------------------------------- | |
| with gr.Blocks(title="SAMatcher") as demo: | |
| gr.Markdown( | |
| """ | |
| # SAMatcher: Dense Co-Visibility Modeling via Cross-View Fusion for Scale-Imbalance Image Matching | |
| Upload (or pick) a pair of images — ideally two views with a strong **scale imbalance** — and | |
| SAMatcher predicts, in a single forward pass, the **co-visible region** shared by both views: | |
| a dense mask and a bounding box per view, plus the cropped matching windows a downstream | |
| matcher would run on. | |
| Model: [`SSSSphinx/SAMatcher`](https://huggingface.co/SSSSphinx/SAMatcher) · | |
| Paper: [SAMatcher](https://huggingface.co/papers/2606.03406) · | |
| Code: [TwSphinx54/SAMatcher](https://github.com/TwSphinx54/SAMatcher) | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| image0_in = gr.Image(label="Image A (reference view)", type="numpy", height=320) | |
| with gr.Column(): | |
| image1_in = gr.Image(label="Image B (scale-imbalanced view)", type="numpy", height=320) | |
| with gr.Row(): | |
| run_btn = gr.Button("Find co-visible region", variant="primary") | |
| with gr.Accordion("Advanced options", open=False): | |
| mask_thresh = gr.Slider(0.1, 0.9, value=0.3, step=0.05, | |
| label="Co-visibility mask threshold") | |
| show_boxes = gr.Checkbox(value=True, label="Draw predicted bounding boxes") | |
| with gr.Row(): | |
| with gr.Column(): | |
| out0 = gr.Image(label="Image A — predicted co-visible mask + box", type="numpy", height=340) | |
| with gr.Column(): | |
| out1 = gr.Image(label="Image B — predicted co-visible mask + box", type="numpy", height=340) | |
| with gr.Row(): | |
| with gr.Column(): | |
| outcrop0 = gr.Image(label="Matching window A", type="numpy", height=260) | |
| with gr.Column(): | |
| outcrop1 = gr.Image(label="Matching window B", type="numpy", height=260) | |
| info_out = gr.Markdown() | |
| run_btn.click( | |
| fn=match, | |
| inputs=[image0_in, image1_in, mask_thresh, show_boxes], | |
| outputs=[out0, out1, outcrop0, outcrop1, info_out], | |
| api_name="match", | |
| concurrency_limit=1, | |
| ) | |
| gr.Examples( | |
| examples=[ | |
| ["assets/sacre_coeur_A.jpg", "assets/sacre_coeur_B.jpg", 0.3, True], | |
| ["assets/toronto_A.jpg", "assets/toronto_B.jpg", 0.3, True], | |
| ["assets/desert_dunes.jpg", "assets/desert_dunes_closeup.jpg", 0.3, True], | |
| ], | |
| inputs=[image0_in, image1_in, mask_thresh, show_boxes], | |
| fn=match, | |
| outputs=[out0, out1, outcrop0, outcrop1, info_out], | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| label="Example pairs (the last one is a truly scale-imbalanced pair)", | |
| ) | |
| gr.Markdown( | |
| """ | |
| --- | |
| **Method**: SAMatcher fuses the two views inside a SAM2.1-Hiera-L encoder with a symmetric | |
| cross-view fusion module (view-aware RoPE) and decodes a shared learnable prompt into | |
| per-view co-visibility masks and boxes — no matching, no correspondence supervision needed. | |
| Example images: Sacré-Cœur & Toronto pairs from [RoMa](https://github.com/Parskatt/RoMa) | |
| (MIT), desert pair from [linoyts/repo-to-space-example-inputs](https://huggingface.co/datasets/linoyts/repo-to-space-example-inputs) (CC0). | |
| """ | |
| ) | |
| custom_css = """ | |
| .dark .gradio-container { | |
| background: var(--background-fill-primary); | |
| color: var(--body-text-color); | |
| } | |
| """ | |
| if __name__ == "__main__": | |
| demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=custom_css) |