add demo

app.py

@@ -0,0 +1,76 @@
+from argparse import ArgumentParser
+from typing import Dict
+import torch
+from PIL import Image
+import modules.transforms as transforms
+from modules.primaps import PriMaPs
+from modules.backbone.dino.dinovit import DinoFeaturizerv2
+from modules.visualization import visualize_demo
+import gradio as gr
+# set seeds
+torch.manual_seed(0)
+torch.cuda.manual_seed(0)
+    
+    
+    
+def gradio_primaps(image_path, threshold, architecture):
+    '''
+    Gradio demo to visualize PriMaPs for a single image.
+    '''    
+    
+    device='cuda:0'   
+    resize_to = 320 if 'v2' not in architecture else 322
+    patch_size = 8 if 'v2' not in architecture else 14
+    
+    # get SLL image encoder and primaps module
+    net = DinoFeaturizerv2(architecture, patch_size)
+    net.to(device)
+    primaps_module = PriMaPs(threshold=threshold,
+                             ignore_id=255)
+    
+    # get transforms
+    demo_transforms = transforms.Compose([transforms.ToTensor(),
+                    transforms.Resize(resize_to), 
+                    transforms.CenterCrop([resize_to, resize_to]),
+                    transforms.Normalize()])
+    
+    # load image and apply transforms
+    image = Image.open(image_path)
+    image, _ = demo_transforms(image, torch.zeros(image.size))
+    image.to(device)
+    # get SSL features
+    feats = net(image.unsqueeze(0).to(device), n=1).squeeze()
+    # get primaps pseudo labels
+    primaps = primaps_module._get_pseudo(image, feats, torch.zeros(image.shape[1:]))
+    # visualize overlay
+    return visualize_demo(image, primaps)
+        
+
+if __name__ == '__main__':    
+    # Example image paths
+    example_images = [
+        "assets/demo_examples/cityscapes_example.png",
+        "assets/demo_examples/coco_example.jpg",
+        "assets/demo_examples/potsdam_example.png"
+    ]
+    
+    # Gradio interface
+    interface = gr.Interface(
+        fn=gradio_primaps,
+        inputs=[
+            gr.Image(type="filepath", label="Image"),
+            gr.Slider(0.0, 1.0, step=0.05, value=0.35, label="Threshold"),
+            gr.Dropdown(choices=['dino_vits', 'dino_vitb', 'dinov2_vits', 'dinov2_vitb'], value='dino_vitb', label="SSL Features"),
+        ],
+        outputs=gr.Image(label="PriMaPs"),
+        title="PriMaPs Demo",
+        description="Upload an image and adjust the threshold to visualize PriMaPs.",
+        examples=[
+            [example_images[0], 0.35, 'dino_vitb'],
+            [example_images[1], 0.35, 'dino_vitb'],
+            [example_images[2], 0.35, 'dino_vitb']
+        ]
+    )
+
+    # Launch the app
+    interface.launch(debug=True)