File size: 7,304 Bytes
159cb1e cf31499 ee88db6 bf4bb3f 159cb1e cf31499 159cb1e ee88db6 159cb1e ee88db6 bf4bb3f ee88db6 bf4bb3f 159cb1e ee88db6 159cb1e ee88db6 159cb1e ee88db6 159cb1e ee88db6 159cb1e 445b4d2 159cb1e 445b4d2 159cb1e ee88db6 159cb1e bf4bb3f 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 194 195 196 197 198 199 200 201 202 203 204 205 206 |
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
import matplotlib.cm as cm
import time
@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", device_map="auto")
inputs = processor(images, return_tensors="pt")
inputs = inputs.to(model.device)
print(
"Model is on device: ",
model.device,
"and inputs are on device: ",
inputs["pixel_values"].device,
)
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))
fig = go.Figure()
# Create colormap (red-yellow-green: red for low scores, green for high scores)
colormap = cm.RdYlGn
# 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()
rgba_color = colormap(color_val)
# Convert to rgba string with transparency
color = f"rgba({int(rgba_color[0] * 255)}, {int(rgba_color[1] * 255)}, {int(rgba_color[2] * 255)}, 0.8)"
hover_text = (
f"Score: {matching_score.item():.3f}<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(
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=0, b=0),
autosize=True,
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 below
2. Click the 'Match Images' button
3. View the matched output image below. Higher scores are green, lower scores are red.
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", scale=2)
# 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()
|