update

app.py

@@ -6,6 +6,7 @@ import torch
 import plotly.graph_objects as go
 from PIL import Image
 import spaces
+import matplotlib.cm as cm
 
 @spaces.GPU
 def process_images(image1, image2):
@@ -17,9 +18,15 @@ def process_images(image1, image2):
 
     images = [image1, image2]
     processor = AutoImageProcessor.from_pretrained("ETH-CVG/lightglue_superpoint")
-    model = AutoModel.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.device,
+    )
     with torch.no_grad():
         outputs = model(**inputs)
 
@@ -39,9 +46,11 @@ def process_images(image1, image2):
     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()
 
+    # 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)
@@ -55,10 +64,13 @@ def process_images(image1, image2):
         output["matching_scores"],
     ):
         color_val = matching_score.item()
-        color = f"rgba({int(255 * (1 - color_val))}, {int(255 * color_val)}, 0, 0.7)"
+        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():.2f}<br>"
+            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})"
         )
@@ -78,7 +90,6 @@ def process_images(image1, image2):
 
     # Update layout to use images as background
     fig.update_layout(
-        title="LightGlue Keypoint Matching",
         xaxis=dict(
             range=[0, width0 + width1],
             showgrid=False,
@@ -93,9 +104,8 @@ def process_images(image1, image2):
             scaleanchor="x",
             scaleratio=1,
         ),
-        margin=dict(l=0, r=0, t=50, b=0),
-        height=max(height0, height1),
-        width=width0 + width1,
+        margin=dict(l=0, r=0, t=0, b=0),
+        autosize=True,
         images=[
             dict(
                 source=pil_img,
@@ -152,7 +162,7 @@ with gr.Blocks(title="LightGlue Matching Demo") as demo:
     process_btn = gr.Button("Match Images", variant="primary")
 
     # Output plot
-    output_plot = gr.Plot(label="Matching Results")
+    output_plot = gr.Plot(label="Matching Results", scale=2)
 
     # Connect the function
     process_btn.click(fn=process_images, inputs=[image1, image2], outputs=output_plot)

requirements.txt

@@ -5,4 +5,5 @@ transformers @ git+https://github.com/huggingface/transformers.git@e5a9ce48f711b
 matplotlib
 torch
 plotly
-spaces
+spaces
+accelerate