DINOv3 satellite roof segmentation app

README.md

@@ -1,13 +1,17 @@
 ---
-title: RoofSegmentation2
-emoji: 💻
-colorFrom: gray
-colorTo: red
+title: Roof Segmentation DINOv3
+emoji: 🛰️
+colorFrom: blue
+colorTo: green
 sdk: gradio
-sdk_version: 6.2.0
+sdk_version: 4.44.0
 app_file: app.py
 pinned: false
-short_description: segmentation using dinov3
+license: other
+models:
+  - facebook/dinov3-vitl16-pretrain-sat493m
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Roof Segmentation with DINOv3 Satellite
+
+Segment roofs from satellite imagery using Meta's DINOv3 model pretrained on 493M satellite images.

app.py

@@ -0,0 +1,188 @@
+import gradio as gr
+import torch
+import torch.nn.functional as F
+import numpy as np
+from PIL import Image
+from transformers import AutoImageProcessor, AutoModel
+from sklearn.cluster import KMeans
+import warnings
+warnings.filterwarnings("ignore")
+
+# Model selection - ViT-L for satellite imagery
+MODEL_NAME = "facebook/dinov3-vitl16-pretrain-sat493m"
+
+print(f"Loading {MODEL_NAME}...")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
+model = AutoModel.from_pretrained(MODEL_NAME).to(device)
+model.eval()
+print(f"Model loaded on {device}")
+
+def extract_features(image):
+    """Extract dense patch features from DINOv3."""
+    inputs = processor(images=image, return_tensors="pt").to(device)
+    
+    with torch.inference_mode():
+        outputs = model(**inputs)
+        # DINOv3: 1 CLS + 4 register tokens + N patch tokens
+        # Skip first 5 tokens (CLS + 4 registers)
+        patch_features = outputs.last_hidden_state[:, 5:, :]
+    
+    return patch_features
+
+def segment_roof(image, num_segments=5, selected_clusters="0"):
+    """
+    Segment roofs using DINOv3 satellite features + K-means.
+    
+    Args:
+        image: Input satellite image
+        num_segments: Number of K-means clusters
+        selected_clusters: Comma-separated cluster indices to highlight as roof
+    """
+    if image is None:
+        return None, None, "Please upload an image"
+    
+    # Convert to PIL if needed
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image).convert("RGB")
+    
+    original_size = image.size  # (W, H)
+    
+    # Extract DINOv3 features
+    features = extract_features(image)
+    
+    # Calculate spatial dimensions
+    # DINOv3 uses patch_size=16
+    num_patches = features.shape[1]
+    h = w = int(np.sqrt(num_patches))
+    
+    # Reshape for clustering
+    feat_np = features.squeeze(0).cpu().numpy()  # [num_patches, hidden_dim]
+    
+    # PCA for dimensionality reduction (helps clustering)
+    from sklearn.decomposition import PCA
+    pca = PCA(n_components=64, random_state=42)
+    feat_reduced = pca.fit_transform(feat_np)
+    
+    # K-means clustering
+    kmeans = KMeans(n_clusters=num_segments, random_state=42, n_init=10)
+    cluster_labels = kmeans.fit_predict(feat_reduced)
+    
+    # Reshape to spatial grid
+    seg_map = cluster_labels.reshape(h, w)
+    
+    # Upscale to original image size
+    seg_resized = np.array(
+        Image.fromarray(seg_map.astype(np.uint8)).resize(
+            original_size, resample=Image.NEAREST
+        )
+    )
+    
+    # Color palette for visualization
+    colors = np.array([
+        [230, 25, 75],    # Red
+        [60, 180, 75],    # Green
+        [255, 225, 25],   # Yellow
+        [0, 130, 200],    # Blue
+        [245, 130, 48],   # Orange
+        [145, 30, 180],   # Purple
+        [70, 240, 240],   # Cyan
+        [240, 50, 230],   # Magenta
+        [210, 245, 60],   # Lime
+        [250, 190, 212],  # Pink
+    ])
+    
+    # Create colored segmentation
+    colored_seg = colors[seg_resized % len(colors)]
+    
+    # Parse selected clusters for roof mask
+    try:
+        roof_indices = [int(x.strip()) for x in selected_clusters.split(",") if x.strip()]
+    except:
+        roof_indices = [0]
+    
+    # Create binary roof mask
+    roof_mask = np.isin(seg_resized, roof_indices).astype(np.uint8) * 255
+    
+    # Create overlay visualization
+    orig_array = np.array(image).astype(np.float32)
+    overlay = orig_array * 0.4 + colored_seg.astype(np.float32) * 0.6
+    
+    # Highlight selected roof clusters
+    for idx in roof_indices:
+        mask = seg_resized == idx
+        overlay[mask] = orig_array[mask] * 0.3 + np.array([255, 0, 0]) * 0.7
+    
+    # Calculate cluster statistics
+    unique, counts = np.unique(seg_resized, return_counts=True)
+    total_pixels = seg_resized.size
+    stats = "**Cluster Statistics:**\n"
+    for u, c in sorted(zip(unique, counts), key=lambda x: -x[1]):
+        pct = (c / total_pixels) * 100
+        marker = " ← ROOF" if u in roof_indices else ""
+        stats += f"- Cluster {u}: {pct:.1f}%{marker}\n"
+    
+    return overlay.astype(np.uint8), roof_mask, stats
+
+# Gradio Interface
+with gr.Blocks(title="Roof Segmentation - DINOv3 Satellite", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # 🛰️ Roof Segmentation with DINOv3 (Satellite)
+    
+    Using Meta's **DINOv3 ViT-L** pretrained on **493M satellite images** at 0.6m resolution.
+    
+    Upload a satellite/aerial image to detect and segment roof areas.
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_image = gr.Image(type="pil", label="📸 Upload Satellite Image")
+            
+            with gr.Accordion("⚙️ Segmentation Settings", open=True):
+                num_segments = gr.Slider(
+                    minimum=3, maximum=12, value=5, step=1,
+                    label="Number of Segments",
+                    info="More segments = finer detail"
+                )
+                selected_clusters = gr.Textbox(
+                    value="0", 
+                    label="Roof Cluster(s)",
+                    info="Enter cluster numbers separated by commas (e.g., '0,2')",
+                    placeholder="0"
+                )
+            
+            segment_btn = gr.Button("🔍 Segment Roofs", variant="primary", size="lg")
+        
+        with gr.Column(scale=2):
+            with gr.Row():
+                output_overlay = gr.Image(label="Segmentation Overlay")
+                output_mask = gr.Image(label="Roof Mask (Binary)")
+            
+            cluster_stats = gr.Markdown(label="Cluster Info")
+    
+    segment_btn.click(
+        fn=segment_roof,
+        inputs=[input_image, num_segments, selected_clusters],
+        outputs=[output_overlay, output_mask, cluster_stats]
+    )
+    
+    gr.Markdown("""
+    ---
+    ### How to Use
+    1. Upload a satellite or aerial image of buildings
+    2. Click **Segment Roofs** to analyze
+    3. Look at the colored overlay - each color is a different segment
+    4. Find which cluster number(s) correspond to roofs (shown in stats)
+    5. Enter those numbers in **Roof Cluster(s)** and re-run
+    6. Download the binary mask for your workflow
+    
+    ### Tips
+    - **Roofs** often cluster together due to similar materials/colors
+    - Try **5-7 segments** for typical suburban imagery
+    - Multiple buildings? Select multiple clusters: `0,3,5`
+    
+    ---
+    *Powered by [DINOv3](https://github.com/facebookresearch/dinov3) pretrained on SAT-493M*
+    """)
+
+demo.launch()

requirements.txt

@@ -0,0 +1,6 @@
+torch
+transformers>=4.40.0
+gradio>=4.0.0
+Pillow
+numpy
+scikit-learn