Upload 2 files

app.py

@@ -0,0 +1,345 @@
+import os
+from typing import Tuple
+
+import gradio as gr
+import numpy as np
+import cv2
+import torch
+import matplotlib.pyplot as plt
+from matplotlib.figure import Figure
+from numpy import ndarray
+import visualize
+
+CSS = """
+#desc, #desc * {
+    text-align: center !important;
+    justify-content: center !important;
+    align-items: center !important;
+}
+"""
+
+DESCRIPTION = """
+<div align="center">
+<h1><ins>MapGlue</ins> 🗺️</h1>
+<h2>
+    MapGlue: Multimodal Remote Sensing Image Matching
+</h2>
+<p>
+    Advanced feature matching system supporting various image modalities including:<br>
+    SAR-Visible, Map-Visible, Depth-Visible, Infrared-Visible, Day-Night matching
+</p>
+</div>
+"""
+
+examples = [
+    [
+        "assets/day-night/L1.png",
+        "assets/day-night/R1.png",
+    ],
+    [
+        "assets/day-night/L2.png",
+        "assets/day-night/R2.png",
+    ],
+    [
+        "assets/depth-visible/L1.jpg",
+        "assets/depth-visible/R1.jpg",
+    ],
+    [
+        "assets/depth-visible/L2.png",
+        "assets/depth-visible/R2.png",
+    ],
+    [
+        "assets/infrared-visible/L1.png",
+        "assets/infrared-visible/R1.png",
+    ],
+    [
+        "assets/infrared-visible/L2.png",
+        "assets/infrared-visible/R2.png",
+    ],
+    [
+        "assets/map-visible/L1.jpg",
+        "assets/map-visible/R1.jpg",
+    ],
+    [
+        "assets/map-visible/L2.png", 
+        "assets/map-visible/R2.png",
+    ],
+    [
+        "assets/sar-visible/L1.jpg",
+        "assets/sar-visible/R1.jpg",
+    ],
+    [
+        "assets/sar-visible/L2.jpg",
+        "assets/sar-visible/R2.jpg",
+    ],
+    [
+        "assets/sar-visible/L3.png",
+        "assets/sar-visible/R3.png",
+    ],
+]
+
+
+def fig_to_ndarray(fig: Figure) -> ndarray:
+    """Convert matplotlib figure to numpy array."""
+    fig.canvas.draw()
+    w, h = fig.canvas.get_width_height()
+    buffer = fig.canvas.buffer_rgba()
+    out = np.frombuffer(buffer, dtype=np.uint8).reshape(h, w, 4)
+    return out
+
+def load_mapglue_model():
+    """Load the MapGlue TorchScript model."""
+    # device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+    device = 'cpu'
+    model_path = './weights/fastmapglue_model.pt'
+    
+    if not os.path.exists(model_path):
+        raise FileNotFoundError(
+            f"Model file not found: {model_path}\n"
+            f"Please ensure the HF_TOKEN environment variable is set to download the model."
+        )
+    
+    model = torch.jit.load(model_path, map_location=device)
+    model.eval()
+    model.to(device)
+    return model, device
+
+
+def run_mapglue_matching(
+    path0: str,
+    path1: str,
+    model_name: str,
+    num_keypoints: int,
+    ransac_threshold: float,
+) -> Tuple[ndarray, ndarray, ndarray, ndarray]:
+    """
+    Run MapGlue matching on two input images using Homography RANSAC.
+    
+    Args:
+        path0, path1: Paths to input images
+        model_name: Name of the matching model (currently supports FastMapGlue)
+        num_keypoints: Number of keypoints to extract
+        ransac_threshold: RANSAC reprojection threshold
+    
+    Returns:
+        Tuple of (raw_keypoint_fig, raw_matching_fig, ransac_keypoint_fig, ransac_matching_fig)
+    """
+    try:
+        # Load model
+        model, device = load_mapglue_model()
+        
+        # Load and preprocess images
+        image0 = cv2.imread(path0)
+        image1 = cv2.imread(path1)
+        
+        if image0 is None or image1 is None:
+            raise ValueError("Could not load one or both images")
+        
+        # Convert BGR to RGB
+        image0 = cv2.cvtColor(image0, cv2.COLOR_BGR2RGB)
+        image1 = cv2.cvtColor(image1, cv2.COLOR_BGR2RGB)
+        
+        # Convert to torch tensors
+        image0_tensor = torch.from_numpy(image0).to(device)
+        image1_tensor = torch.from_numpy(image1).to(device)
+        num_keypoints_tensor = torch.tensor(num_keypoints).to(device)
+        
+        # Run inference
+        with torch.no_grad():
+            points_tensor = model(image0_tensor, image1_tensor, num_keypoints_tensor)
+            points0 = points_tensor[:, :2]
+            points1 = points_tensor[:, 2:]
+        
+        # Create raw matching visualization
+        plt.figure(figsize=(12, 6))
+        axes = visualize.show_images([image0, image1])
+        visualize.draw_matches(points0, points1, line_colors="lime", line_width=0.8)
+        visualize.add_text(0, f'Raw matches: {len(points0)}', font_size=16)
+        raw_matching_fig = fig_to_ndarray(plt.gcf())
+        
+        # Create raw keypoints visualization
+        plt.figure(figsize=(12, 6))
+        axes = visualize.show_images([image0, image1])
+        visualize.draw_keypoints([points0.cpu().numpy(), points1.cpu().numpy()], 
+                               kp_color=["lime", "lime"], kp_size=20)
+        visualize.add_text(0, f'Raw keypoints: {len(points0)}', font_size=16)
+        raw_keypoint_fig = fig_to_ndarray(plt.gcf())
+        
+        # Apply RANSAC filtering
+        points0_np = points0.cpu().numpy()
+        points1_np = points1.cpu().numpy()
+
+  
+        H_pred, inlier_mask = cv2.findHomography(
+            points0_np, points1_np,
+            cv2.USAC_MAGSAC,
+            ransacReprojThreshold=ransac_threshold,
+            maxIters=10000,
+            confidence=0.9999
+        )
+
+        if inlier_mask is not None and inlier_mask.sum() > 0:
+            inlier_mask = inlier_mask.ravel() > 0
+            mkpts0 = points0_np[inlier_mask]
+            mkpts1 = points1_np[inlier_mask]
+            
+            # Create RANSAC matching visualization
+            plt.figure(figsize=(12, 6))
+            axes = visualize.show_images([image0, image1])
+            visualize.draw_matches(mkpts0, mkpts1, line_colors="lime", line_width=1)
+            visualize.add_text(0, f'RANSAC matches @{ransac_threshold}px: {len(mkpts0)}/{len(points0)}', font_size=16)
+            ransac_matching_fig = fig_to_ndarray(plt.gcf())
+            
+            # Create RANSAC keypoints visualization
+            plt.figure(figsize=(12, 6))
+            axes = visualize.show_images([image0, image1])
+            visualize.draw_keypoints([mkpts0, mkpts1], 
+                                   kp_color=["lime", "lime"], kp_size=20)
+            visualize.add_text(0, f'RANSAC keypoints @{ransac_threshold}px: {len(mkpts0)}', font_size=16)
+            ransac_keypoint_fig = fig_to_ndarray(plt.gcf())
+        else:
+            # No inliers found
+            ransac_matching_fig = None
+            ransac_keypoint_fig = None
+        
+        plt.close('all')  # Clean up matplotlib figures
+        
+        return (
+            raw_keypoint_fig,
+            raw_matching_fig, 
+            ransac_keypoint_fig,
+            ransac_matching_fig,
+        )
+        
+    except Exception as e:
+        print(f"Error in matching: {str(e)}")
+        # Return empty arrays in case of error
+        empty_img = np.zeros((400, 800, 4), dtype=np.uint8)
+        return (empty_img, empty_img, empty_img, empty_img)
+
+
+with gr.Blocks(css=CSS) as demo:
+    with gr.Tab("Image Matching"):
+        with gr.Row():
+            with gr.Column(scale=3):
+                gr.HTML(DESCRIPTION, elem_id="desc")
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("### Input Panels:")
+                with gr.Row():
+                    model_name = gr.Dropdown(
+                        choices=["FastMapGlue"],
+                        value="FastMapGlue",
+                        label="Matching Model",
+                    )
+                with gr.Row():
+                    path0 = gr.Image(
+                        height=300,
+                        image_mode="RGB",
+                        type="filepath",
+                        label="Image 0",
+                    )
+                    path1 = gr.Image(
+                        height=300,
+                        image_mode="RGB", 
+                        type="filepath",
+                        label="Image 1",
+                    )
+                with gr.Row():
+                    stop = gr.Button(value="Stop", variant="stop")
+                    run = gr.Button(value="Run", variant="primary")
+                    
+                with gr.Accordion("Advanced Settings", open=False):
+                    with gr.Accordion("Matching Settings"):
+                        with gr.Row():
+                            num_keypoints = gr.Slider(
+                                minimum=512,
+                                maximum=4096,
+                                value=2048,
+                                step=256,
+                                label="Number of Keypoints",
+                            )
+                    with gr.Accordion("RANSAC Settings"):
+                        with gr.Row():
+                            ransac_threshold = gr.Slider(
+                                minimum=0.5,
+                                maximum=10.0,
+                                value=5.0,
+                                step=0.5,
+                                label="RANSAC Threshold",
+                            )
+                            
+                with gr.Row():
+                    with gr.Accordion("Example Pairs"):
+                        gr.Examples(
+                            examples=examples,
+                            inputs=[path0, path1],
+                            label="Click an example pair below",
+                        )
+                        
+            with gr.Column():
+                gr.Markdown(
+                    "### Output Panels"
+                )
+                with gr.Accordion("Raw Keypoints", open=False):
+                    raw_keypoint_fig = gr.Image(
+                        format="png", type="numpy", label="Raw Keypoints"
+                    )
+                with gr.Accordion("Raw Matches"):
+                    raw_matching_fig = gr.Image(
+                        format="png", type="numpy", label="Raw Matches"
+                    )
+                with gr.Accordion("RANSAC Keypoints", open=False):
+                    ransac_keypoint_fig = gr.Image(
+                        format="png", type="numpy", label="RANSAC Keypoints"
+                    )
+                with gr.Accordion("RANSAC Matches"):
+                    ransac_matching_fig = gr.Image(
+                        format="png", type="numpy", label="RANSAC Matches"
+                    )
+
+        inputs = [
+            path0,
+            path1,
+            model_name,
+            num_keypoints,
+            ransac_threshold,
+        ]
+        outputs = [
+            raw_keypoint_fig,
+            raw_matching_fig,
+            ransac_keypoint_fig,
+            ransac_matching_fig,
+        ]
+
+        running_event = run.click(
+            fn=run_mapglue_matching, inputs=inputs, outputs=outputs
+        )
+        stop.click(
+            fn=None, inputs=None, outputs=None, cancels=[running_event]
+        )
+
+if __name__ == "__main__":
+    # Download model weights on startup if HF_TOKEN is available
+    HF_TOKEN = os.getenv("HF_TOKEN")
+    if HF_TOKEN:
+        model_path = './weights/fastmapglue_model.pt'
+        if not os.path.exists(model_path):
+            try:
+                import requests
+                # 使用 resolve 来直接下载文件
+                model_url = "https://huggingface.co/wupeihao/mapglue/resolve/main/fastmapglue_model.pt"
+                headers = {"Authorization": f"Bearer {HF_TOKEN}"}
+                
+                print("Downloading MapGlue model...")
+                response = requests.get(model_url, headers=headers)
+                response.raise_for_status()
+                
+                os.makedirs('./weights', exist_ok=True)
+                with open(model_path, 'wb') as f:
+                    f.write(response.content)
+                print("Model downloaded successfully!")
+            except Exception as e:
+                print(f"Failed to download model: {str(e)}")
+    
+    demo.launch()

visualize.py

@@ -0,0 +1,156 @@
+import matplotlib
+import matplotlib.patheffects as peffects
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+
+def show_images(image_list, titles=None, colormaps="gray", dpi=100, pad=0.5, auto_size=True):
+    """
+    Display a set of images horizontally.
+
+    Args:
+        image_list: List of images in either NumPy RGB (H, W, 3), 
+                    PyTorch RGB (3, H, W) or grayscale (H, W) format.
+        titles: List of titles for each image.
+        colormaps: Colormap for grayscale images.
+        dpi: Figure resolution.
+        pad: Padding between images.
+        auto_size: Whether the figure size should adapt to the images' aspect ratios.
+    """
+    # Convert torch.Tensor images to NumPy arrays in (H, W, 3) format.
+    image_list = [
+        img.permute(1, 2, 0).cpu().numpy()
+        if (isinstance(img, torch.Tensor) and img.dim() == 3)
+        else img
+        for img in image_list
+    ]
+    num_imgs = len(image_list)
+    if not isinstance(colormaps, (list, tuple)):
+        colormaps = [colormaps] * num_imgs
+
+    if auto_size:
+        ratios = [im.shape[1] / im.shape[0] for im in image_list]  # width / height
+    else:
+        ratios = [4 / 3] * num_imgs
+    fig_size = [sum(ratios) * 4.5, 4.5]
+    fig, axes = plt.subplots(1, num_imgs, figsize=fig_size, dpi=dpi, gridspec_kw={"width_ratios": ratios})
+    if num_imgs == 1:
+        axes = [axes]
+    for i in range(num_imgs):
+        axes[i].imshow(image_list[i], cmap=plt.get_cmap(colormaps[i]))
+        axes[i].set_xticks([])
+        axes[i].set_yticks([])
+        axes[i].set_axis_off()
+        for spine in axes[i].spines.values():
+            spine.set_visible(False)
+        if titles:
+            axes[i].set_title(titles[i])
+    fig.tight_layout(pad=pad)
+
+
+def draw_keypoints(keypoints, kp_color="lime", kp_size=4, ax_list=None, alpha_value=1.0):
+    """
+    Plot keypoints on existing images.
+
+    Args:
+        keypoints: List of ndarrays (N, 2) for each set of keypoints.
+        kp_color: Color for keypoints, or list of colors for each set.
+        kp_size: Size of keypoints.
+        ax_list: List of axes to plot keypoints on; defaults to current figure's axes.
+        alpha_value: Opacity for keypoints.
+    """
+    if not isinstance(kp_color, list):
+        kp_color = [kp_color] * len(keypoints)
+    if not isinstance(alpha_value, list):
+        alpha_value = [alpha_value] * len(keypoints)
+    if ax_list is None:
+        ax_list = plt.gcf().axes
+    for ax, pts, color, alpha in zip(ax_list, keypoints, kp_color, alpha_value):
+        if isinstance(pts, torch.Tensor):
+            pts = pts.cpu().numpy()
+        ax.scatter(pts[:, 0], pts[:, 1], c=color, s=kp_size, linewidths=0, alpha=alpha)
+
+
+def draw_matches(pts_left, pts_right, line_colors=None, line_width=1.5, endpoint_size=4, alpha_value=1.0, labels=None, axes_pair=None):
+    """
+    Draw matches between a pair of images.
+
+    Args:
+        pts_left, pts_right: Corresponding keypoints for the two images (N, 2).
+        line_colors: Colors for each match line, either as a string or an RGB tuple.
+                     If not provided, random colors will be generated.
+        line_width: Width of the match lines.
+        endpoint_size: Size of the endpoints (if 0, endpoints are not drawn).
+        alpha_value: Opacity for the match lines.
+        labels: Optional list of labels for each match.
+        axes_pair: List of two axes [ax_left, ax_right] to plot the images; defaults to the first two axes in the current figure.
+    """
+    fig = plt.gcf()
+    if axes_pair is None:
+        axs = fig.axes
+        ax_left, ax_right = axs[0], axs[1]
+    else:
+        ax_left, ax_right = axes_pair
+    if isinstance(pts_left, torch.Tensor):
+        pts_left = pts_left.cpu().numpy()
+    if isinstance(pts_right, torch.Tensor):
+        pts_right = pts_right.cpu().numpy()
+    assert len(pts_left) == len(pts_right)
+    if line_colors is None:
+        line_colors = matplotlib.cm.hsv(np.random.rand(len(pts_left))).tolist()
+    elif len(line_colors) > 0 and not isinstance(line_colors[0], (tuple, list)):
+        line_colors = [line_colors] * len(pts_left)
+
+    if line_width > 0:
+        for i in range(len(pts_left)):
+            connector = matplotlib.patches.ConnectionPatch(
+                xyA=(pts_left[i, 0], pts_left[i, 1]),
+                xyB=(pts_right[i, 0], pts_right[i, 1]),
+                coordsA=ax_left.transData,
+                coordsB=ax_right.transData,
+                axesA=ax_left,
+                axesB=ax_right,
+                zorder=1,
+                color=line_colors[i],
+                linewidth=line_width,
+                clip_on=True,
+                alpha=alpha_value,
+                label=None if labels is None else labels[i],
+                picker=5.0,
+            )
+            connector.set_annotation_clip(True)
+            fig.add_artist(connector)
+
+    # Freeze axis autoscaling to prevent changes.
+    ax_left.autoscale(enable=False)
+    ax_right.autoscale(enable=False)
+
+    if endpoint_size > 0:
+        ax_left.scatter(pts_left[:, 0], pts_left[:, 1], c=line_colors, s=endpoint_size)
+        ax_right.scatter(pts_right[:, 0], pts_right[:, 1], c=line_colors, s=endpoint_size)
+
+
+def add_text(axis_idx, text, pos=(0.01, 0.99), font_size=15, txt_color="w", border_color="k", border_width=2, h_align="left", v_align="top"):
+    """
+    Add an annotation with an outline to a specified axis.
+
+    Args:
+        axis_idx: Index of the axis in the current figure where the annotation will be added.
+        text: The annotation text.
+        pos: Position of the annotation in axis coordinates (e.g., (0.01, 0.99)).
+        font_size: Font size of the text.
+        txt_color: Text color.
+        border_color: Outline color (if None, no outline is applied).
+        border_width: Width of the outline.
+        h_align: Horizontal alignment (e.g., "left").
+        v_align: Vertical alignment (e.g., "top").
+    """
+    current_ax = plt.gcf().axes[axis_idx]
+    annotation = current_ax.text(
+        *pos, text, fontsize=font_size, ha=h_align, va=v_align, color=txt_color, transform=current_ax.transAxes
+    )
+    if border_color is not None:
+        annotation.set_path_effects([
+            peffects.Stroke(linewidth=border_width, foreground=border_color),
+            peffects.Normal(),
+        ])