File size: 6,890 Bytes
159cb1e cf31499 159cb1e cf31499 159cb1e d8df29d 159cb1e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 |
import gradio as gr
import numpy as np
from transformers import AutoImageProcessor, AutoModel
from transformers.image_utils import to_numpy_array
import torch
import plotly.graph_objects as go
from PIL import Image
import spaces
@spaces.GPU
def process_images(image1, image2):
"""
Process two images and return a plot of the matching keypoints.
"""
if image1 is None or image2 is None:
return None
images = [image1, image2]
processor = AutoImageProcessor.from_pretrained("ETH-CVG/lightglue_superpoint")
model = AutoModel.from_pretrained("ETH-CVG/lightglue_superpoint")
inputs = processor(images, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
image_sizes = [[(image.height, image.width) for image in images]]
outputs = processor.post_process_keypoint_matching(
outputs, image_sizes, threshold=0.2
)
output = outputs[0]
image1 = to_numpy_array(image1)
image2 = to_numpy_array(image2)
height0, width0 = image1.shape[:2]
height1, width1 = image2.shape[:2]
# Create PIL image from numpy array
pil_img = Image.fromarray((image1 / 255.0 * 255).astype(np.uint8))
pil_img2 = Image.fromarray((image2 / 255.0 * 255).astype(np.uint8))
# Create Plotly figure
fig = go.Figure()
# Get keypoints
keypoints0_x, keypoints0_y = output["keypoints0"].unbind(1)
keypoints1_x, keypoints1_y = output["keypoints1"].unbind(1)
# Add a separate trace for each match (line + markers) to enable highlighting
for keypoint0_x, keypoint0_y, keypoint1_x, keypoint1_y, matching_score in zip(
keypoints0_x,
keypoints0_y,
keypoints1_x,
keypoints1_y,
output["matching_scores"],
):
color_val = matching_score.item()
color = f"rgba({int(255 * (1 - color_val))}, {int(255 * color_val)}, 0, 0.7)"
hover_text = (
f"Score: {matching_score.item():.2f}<br>"
f"Point 1: ({keypoint0_x.item():.1f}, {keypoint0_y.item():.1f})<br>"
f"Point 2: ({keypoint1_x.item():.1f}, {keypoint1_y.item():.1f})"
)
fig.add_trace(
go.Scatter(
x=[keypoint0_x.item(), keypoint1_x.item() + width0],
y=[keypoint0_y.item(), keypoint1_y.item()],
mode="lines+markers",
line=dict(color=color, width=2),
marker=dict(color=color, size=5, opacity=0.8),
hoverinfo="text",
hovertext=hover_text,
showlegend=False,
)
)
# Update layout to use images as background
fig.update_layout(
title="LightGlue Keypoint Matching",
xaxis=dict(
range=[0, width0 + width1],
showgrid=False,
zeroline=False,
showticklabels=False,
),
yaxis=dict(
range=[max(height0, height1), 0],
showgrid=False,
zeroline=False,
showticklabels=False,
scaleanchor="x",
scaleratio=1,
),
margin=dict(l=0, r=0, t=50, b=0),
height=max(height0, height1),
width=width0 + width1,
images=[
dict(
source=pil_img,
xref="x",
yref="y",
x=0,
y=0,
sizex=width0,
sizey=height0,
sizing="stretch",
opacity=1,
layer="below",
),
dict(
source=pil_img2,
xref="x",
yref="y",
x=width0,
y=0,
sizex=width1,
sizey=height1,
sizing="stretch",
opacity=1,
layer="below",
),
],
)
return fig
# Create the Gradio interface
with gr.Blocks(title="LightGlue Matching Demo") as demo:
gr.Markdown("# LightGlue Matching Demo")
gr.Markdown(
"Upload two images and get a side-by-side matching of your images using LightGlue."
)
gr.Markdown("""
## How to use:
1. Upload two images using the file uploaders above
2. Click the 'Match Images' button
3. View the matched output image below
The app will create a side-by-side matching of your images using LightGlue.
You can also select an example image pair from the dataset.
""")
with gr.Row():
# Input images on the same row
image1 = gr.Image(label="First Image", type="pil")
image2 = gr.Image(label="Second Image", type="pil")
# Process button
process_btn = gr.Button("Match Images", variant="primary")
# Output plot
output_plot = gr.Plot(label="Matching Results")
# Connect the function
process_btn.click(fn=process_images, inputs=[image1, image2], outputs=output_plot)
# Add some example usage
examples = gr.Dataset(
components=[image1, image2],
label="Example Image Pairs",
samples=[
[
"https://raw.githubusercontent.com/magicleap/SuperGluePretrainedNetwork/refs/heads/master/assets/phototourism_sample_images/united_states_capitol_98169888_3347710852.jpg",
"https://raw.githubusercontent.com/magicleap/SuperGluePretrainedNetwork/refs/heads/master/assets/phototourism_sample_images/united_states_capitol_26757027_6717084061.jpg",
],
[
"https://raw.githubusercontent.com/cvg/LightGlue/refs/heads/main/assets/DSC_0410.JPG",
"https://raw.githubusercontent.com/cvg/LightGlue/refs/heads/main/assets/DSC_0411.JPG",
],
[
"https://raw.githubusercontent.com/cvg/LightGlue/refs/heads/main/assets/sacre_coeur1.jpg",
"https://raw.githubusercontent.com/cvg/LightGlue/refs/heads/main/assets/sacre_coeur2.jpg",
],
[
"https://raw.githubusercontent.com/magicleap/SuperGluePretrainedNetwork/refs/heads/master/assets/phototourism_sample_images/piazza_san_marco_06795901_3725050516.jpg",
"https://raw.githubusercontent.com/magicleap/SuperGluePretrainedNetwork/refs/heads/master/assets/phototourism_sample_images/piazza_san_marco_58751010_4849458397.jpg",
],
[
"https://raw.githubusercontent.com/magicleap/SuperGluePretrainedNetwork/refs/heads/master/assets/phototourism_sample_images/london_bridge_19481797_2295892421.jpg",
"https://raw.githubusercontent.com/magicleap/SuperGluePretrainedNetwork/refs/heads/master/assets/phototourism_sample_images/london_bridge_78916675_4568141288.jpg",
],
],
)
examples.select(lambda x: (x[0], x[1]), [examples], [image1, image2])
if __name__ == "__main__":
demo.launch()
|