Update app.py

app.py

@@ -14,8 +14,8 @@ from transformers import AutoImageProcessor, AutoModel
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 MODEL_MAP = {
-    "DINOv3 ViT-L/16 Satellite": "facebook/dinov3-vitl16-pretrain-sat493m",
-    "DINOv3 ViT-L/16 LVD (web data)": "facebook/dinov3-vitl16-pretrain-lvd1689m",
+    "DINOv3 ViT-L/16 Satellite (493M)": "facebook/dinov3-vitl16-pretrain-sat493m",
+    "DINOv3 ViT-L/16 LVD (1.7B web)": "facebook/dinov3-vitl16-pretrain-lvd1689m",
     "⚠️ DINOv3 ViT-7B/16 Satellite": "facebook/dinov3-vit7b16-pretrain-sat493m",
 }
 
@@ -27,40 +27,54 @@ model = None
 
 
 def cleanup_memory():
+    """Aggressive memory cleanup for model switching"""
+    global processor, model
+
+    if model is not None:
+        del model
+        model = None
+
+    if processor is not None:
+        del processor
+        processor = None
 
     gc.collect()
 
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
+        torch.cuda.synchronize()
 
 
 def load_model(name):
     """Load model with proper memory management and dtype handling"""
     global processor, model
 
-    # Clean up existing model
-    del model
-    cleanup_memory()
+    try:
+        # Clean up existing model
+        cleanup_memory()
 
-    model_id = MODEL_MAP[name]
+        model_id = MODEL_MAP[name]
 
-    # Load processor
-    processor = AutoImageProcessor.from_pretrained(model_id)
+        # Load processor
+        processor = AutoImageProcessor.from_pretrained(model_id)
 
-    model = (
-        AutoModel.from_pretrained(
-            model_id,
-            torch_dtype="auto",
+        model = (
+            AutoModel.from_pretrained(
+                model_id,
+                torch_dtype="auto",
+            )
+            .to(DEVICE)
+            .eval()
         )
-        .to(DEVICE)
-        .eval()
-    )
 
-    # Get model info
-    param_count = sum(p.numel() for p in model.parameters()) / 1e9
+        # Get model info
+        param_count = sum(p.numel() for p in model.parameters()) / 1e9
 
-    return f"Loaded: {name} | {param_count:.1f}B params | {DEVICE.upper()}"
+        return f"✅ Loaded: {name} | {param_count:.1f}B params | {DEVICE.upper()}"
 
+    except Exception as e:
+        cleanup_memory()
+        return f"❌ Failed to load {name}: {str(e)}"
 
 
 # Initialize default model
@@ -93,7 +107,8 @@ def _overlay(orig, heat01, alpha=0.55, box=None):
     """Create heatmap overlay"""
     H, W = orig.height, orig.width
     heat = Image.fromarray((heat01 * 255).astype(np.uint8)).resize((W, H))
-    rgba = (cm.get_cmap("inferno")(np.asarray(heat) / 255.0) * 255).astype(np.uint8)
+    # Use turbo colormap - better for satellite imagery
+    rgba = (cm.get_cmap("turbo")(np.asarray(heat) / 255.0) * 255).astype(np.uint8)
     ov = Image.fromarray(rgba, "RGBA")
     ov.putalpha(int(alpha * 255))
     base = orig.copy().convert("RGBA")
@@ -101,8 +116,13 @@ def _overlay(orig, heat01, alpha=0.55, box=None):
     if box:
         from PIL import ImageDraw
 
-        ImageDraw.Draw(out, "RGBA").rectangle(
-            box, outline=(255, 255, 255, 220), width=2
+        draw = ImageDraw.Draw(out, "RGBA")
+        # Enhanced box visualization
+        draw.rectangle(box, outline=(255, 255, 255, 255), width=3)
+        draw.rectangle(
+            (box[0] - 1, box[1] - 1, box[2] + 1, box[3] + 1),
+            outline=(0, 0, 0, 200),
+            width=1,
         )
     return out
 
@@ -153,34 +173,97 @@ def reset():
     return None, None
 
 
-with gr.Blocks() as demo:
-    gr.Markdown("## DINOv3: patch similarity visualizer")
-    gr.Markdown(
-        "This is an app where you can upload an image, click on an object in the image and get the most similar patches to it according to DINOv3, revealing the way DINOv3 segments objects through features natively."
-    )
-    gr.Markdown("There's multiple model options you can pick from the dropdown.")
-    gr.Markdown(
-        "Please click Reset before you want to upload another image, as this app keeps features of the images."
+with gr.Blocks(
+    theme=gr.themes.Citrus(),
+    css="""
+    .container {max-width: 1200px; margin: auto;}
+    .header {text-align: center; padding: 20px;}
+    .info-box {
+        background: rgba(0,0,0,0.03); 
+        border-radius: 8px; 
+        padding: 12px; 
+        margin: 10px 0;
+        border-left: 4px solid #2563eb;
+    }
+    """,
+) as demo:
+    gr.HTML(
+        """
+    <div class="header">
+        <h1>🛰️ DINOv3 Satellite Vision: Interactive Patch Similarity</h1>
+        <p style="font-size: 1.1em; color: #666;">
+            Click any region to visualize feature similarities across the image
+        </p>
+    </div>
+    """
     )
 
-    with gr.Column():
-        with gr.Row(scale=1):
+    with gr.Row():
+        with gr.Column(scale=3):
+            gr.Markdown(
+                """
+            ### How it works
+            1. **Select a model** - Satellite-pretrained models optimized for aerial/satellite imagery
+            2. **Upload or select an image** - Works best with satellite, aerial, or outdoor scenes
+            3. **Click any region** - See how similar other patches are to your selection
+            4. **Adjust opacity** - Fine-tune visualization clarity
+            """
+            )
+
+        with gr.Column(scale=2):
+            gr.HTML(
+                """
+            <div class="info-box">
+                <b>💡 Model Info:</b><br>
+                • <b>Satellite (493M)</b>: Trained on 493M satellite images<br>
+                • <b>LVD (1.7B)</b>: Trained on 1.7B diverse web images<br>
+                • <b>7B Satellite</b>: Massive capacity, requires high VRAM
+            </div>
+            """
+            )
+
+    with gr.Row():
+        with gr.Column(scale=1):
             model_choice = gr.Dropdown(
-                choices=list(MODEL_MAP.keys()), value=DEFAULT_NAME, label="Model"
+                choices=list(MODEL_MAP.keys()),
+                value=DEFAULT_NAME,
+                label="🤖 Model Selection",
+                info="Satellite models excel at geographic and structural patterns",
             )
             status = gr.Textbox(
-                label="Status", value=f"Loaded: {DEFAULT_NAME}", interactive=False
+                label="📡 Model Status",
+                value=f"✅ Loaded: {DEFAULT_NAME}",
+                interactive=False,
+                lines=1,
             )
-            opacity = gr.Slider(0.0, 1.0, 0.55, step=0.05, label="Opacity for the Map")
+            opacity = gr.Slider(
+                0.0,
+                1.0,
+                0.55,
+                step=0.05,
+                label="🎨 Heatmap Opacity",
+                info="Balance between image and similarity map",
+            )
+
+            with gr.Row():
+                reset_btn = gr.Button("🔄 Reset", variant="secondary", scale=1)
+                clear_btn = gr.ClearButton(value="🗑️ Clear All", scale=1)
 
-        with gr.Row(scale=1):
+        with gr.Column(scale=2):
             img = gr.Image(
-                type="pil", label="Image", interactive=True, height=750, width=750
+                type="pil",
+                label="Interactive Canvas (Click to explore)",
+                interactive=True,
+                height=600,
+                show_download_button=True,
+                show_share_button=False,
             )
 
     state = gr.State()
 
-    model_choice.change(load_model, inputs=model_choice, outputs=status)
+    model_choice.change(
+        load_model, inputs=model_choice, outputs=status, show_progress="full"
+    )
 
     img.upload(prepare, inputs=img, outputs=state)
 
@@ -191,11 +274,12 @@ with gr.Blocks() as demo:
         show_progress="minimal",
     )
 
-    gr.Button("Reset").click(reset, outputs=[img, state])
+    reset_btn.click(reset, outputs=[img, state])
+    clear_btn.add([img, state])
 
     # Examples from current directory
     example_files = [
-        str(f)
+        f.name
         for f in Path.cwd().iterdir()
         if f.suffix.lower() in [".jpg", ".jpeg", ".png", ".webp"]
     ]
@@ -206,10 +290,22 @@ with gr.Blocks() as demo:
             inputs=img,
             fn=prepare,
             outputs=[state],
-            label="Try an example image and then click on the objects.",
+            label="Example Images",
             examples_per_page=4,
             cache_examples=False,
         )
 
+    gr.Markdown(
+        """
+    ---
+    <div style="text-align: center; color: #666; font-size: 0.9em;">
+        <b>Performance Notes:</b> Satellite models are optimized for geographic patterns, land use classification, 
+        and structural analysis. The 7B model provides exceptional detail but requires significant compute.
+        <br><br>
+        Built with DINOv3 | Optimized for satellite and aerial imagery analysis
+    </div>
+    """
+    )
+
 if __name__ == "__main__":
     demo.launch(share=False, debug=True)