feat: added disk variant to the space

app.py

@@ -11,7 +11,7 @@ import time
 
 
 @spaces.GPU
-def process_images(image1, image2):
+def process_images(image1, image2, detector_choice):
     """
     Process two images and return a plot of the matching keypoints.
     """
@@ -19,8 +19,17 @@ def process_images(image1, image2):
         return None
 
     images = [image1, image2]
-    processor = AutoImageProcessor.from_pretrained("ETH-CVG/lightglue_superpoint")
-    model = AutoModel.from_pretrained("ETH-CVG/lightglue_superpoint", device_map="auto")
+    
+    # Select model based on detector choice
+    if detector_choice == "SuperPoint":
+        model_name = "ETH-CVG/lightglue_superpoint"
+        trust_remote_code = False
+    else:  # DISK
+        model_name = "ETH-CVG/lightglue_disk"
+        trust_remote_code = True
+    
+    processor = AutoImageProcessor.from_pretrained(model_name, trust_remote_code=trust_remote_code)
+    model = AutoModel.from_pretrained(model_name, device_map="auto", trust_remote_code=trust_remote_code)
     inputs = processor(images, return_tensors="pt")
     inputs = inputs.to(model.device)
     print(
@@ -148,14 +157,24 @@ with gr.Blocks(title="LightGlue Matching Demo") as demo:
     )
     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.
+        1. Select a detector (SuperPoint or DISK)
+        2. Upload two images using the file uploaders below
+        3. Click the 'Match Images' button
+        4. 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():
+        # Detector choice selector
+        detector_choice = gr.Radio(
+            choices=["SuperPoint", "DISK"],
+            value="SuperPoint",
+            label="Detector Choice",
+            info="Choose between SuperPoint or DISK detector"
+        )
+
     with gr.Row():
         # Input images on the same row
         image1 = gr.Image(label="First Image", type="pil")
@@ -168,7 +187,7 @@ with gr.Blocks(title="LightGlue Matching Demo") as demo:
     output_plot = gr.Plot(label="Matching Results", scale=2)
 
     # Connect the function
-    process_btn.click(fn=process_images, inputs=[image1, image2], outputs=[output_plot])
+    process_btn.click(fn=process_images, inputs=[image1, image2, detector_choice], outputs=[output_plot])
 
     # Add some example usage
 

requirements.txt

@@ -1,9 +1,10 @@
 gradio>=5.34.2
 Pillow>=10.0.0
 numpy>=1.24.0 
-transformers @ git+https://github.com/huggingface/transformers.git@e5a9ce48f711b1f26eef3f7047a13b8235e4a71b
+transformers @ git+https://github.com/huggingface/transformers.git@1255480fd226129075e10c20842efd444f5b0e36
 matplotlib
 torch
 plotly
 spaces
-accelerate
+accelerate
+kornia