Initial Space upload

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/example.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,12 +1,18 @@
----
-title: Road Detection Model
-emoji: 🌖
-colorFrom: purple
-colorTo: yellow
-sdk: gradio
-sdk_version: 5.43.1
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# PyTorch `.pth` Model – Hugging Face Space
+
+This Space hosts a custom PyTorch model loaded from `weights/road_detection_model.pth`.
+
+## Inputs
+- Satellite Image
+
+## Programmatic Usage
+
+```python
+from gradio_client import Client
+
+client = Client("https://<ORG_OR_USER>-<SPACE_NAME>.hf.space")
+res = client.predict(
+    {x: '<IMAGE>'},   
+    api_name="/predict"
+)
+print(res)  # {"pred": int, "probs": [ ... ]}

app.py

@@ -0,0 +1,319 @@
+import os
+import json
+import torch
+import gradio as gr
+import cv2
+import numpy as np
+from PIL import Image
+from torchvision import transforms
+import base64
+from io import BytesIO
+
+from model_def import build_model
+
+# ---------- Config ----------
+WEIGHTS_PATH = os.environ.get("WEIGHTS_PATH", "weights/road_detection_model.pth")
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+MODEL = None
+MODEL_EVAL = True  # turn off dropout/bn if True
+
+THRESHOLD = 0.9
+MIN_AREA = 25
+KERNEL_SIZE = 3
+ASPECT_RATIO = 2
+PERIMETER = 1
+CONNECTIVITY = 8
+
+# If your model expects different input (e.g., images, text), adapt preprocess/postprocess accordingly.
+def load_model():
+    global MODEL
+    MODEL = build_model()
+    # If you saved state_dict:
+    state = torch.load(WEIGHTS_PATH, map_location="cpu")
+    MODEL.load_state_dict(state)
+    MODEL.to(DEVICE)
+    if MODEL_EVAL:
+        MODEL.eval()
+
+def preprocess(raw_input):
+    """
+    Preprocess the input image for road detection.
+    
+    Args:
+        raw_input: Can be base64 string, file path, or PIL Image
+        
+    Returns:
+        torch.Tensor: Preprocessed image tensor ready for model inference
+    """
+    # Handle different input types
+    if isinstance(raw_input, str):
+        # Check if it's a base64 string
+        if raw_input.startswith('data:image'):
+            # Extract base64 data
+            base64_data = raw_input.split(',')[1]
+            image_data = base64.b64decode(base64_data)
+            image = Image.open(BytesIO(image_data)).convert('RGB')
+        else:
+            # Assume it's a file path
+            image = Image.open(raw_input).convert('RGB')
+    elif isinstance(raw_input, dict) and 'image' in raw_input:
+        # Handle Gradio image input
+        image = Image.fromarray(raw_input['image']).convert('RGB')
+    else:
+        # Assume it's already a PIL Image
+        image = raw_input.convert('RGB')
+    
+    # Resize image to model input size (256x256)
+    input_size = 256
+    image = image.resize((input_size, input_size), Image.LANCZOS)
+    
+    # Convert to tensor and normalize
+    transform = transforms.Compose([
+        transforms.ToTensor(),
+    ])
+    
+    # Add batch dimension
+    tensor = transform(image).unsqueeze(0)
+    
+    return tensor.to(DEVICE)
+
+def _clean_road_mask(mask: np.ndarray, min_area: int = MIN_AREA, kernel_size: int = KERNEL_SIZE) -> np.ndarray:
+    """
+    Clean the road mask by removing small disconnected segments and improving connectivity.
+    Preserves long, thin roads by considering aspect ratio and perimeter.
+    
+    Args:
+        mask (np.ndarray): Binary road mask (0-255)
+        min_area (int): Minimum area in pixels for a road segment to be kept
+        kernel_size (int): Size of morphological operation kernel
+        
+    Returns:
+        np.ndarray: Cleaned binary road mask
+    """
+    # Convert to binary (0 or 1)
+    binary_mask = (mask > 127).astype(np.uint8)
+    
+    # Create kernel for morphological operations
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
+    
+    # 1. Remove small noise with opening (erosion followed by dilation)
+    cleaned = cv2.morphologyEx(binary_mask, cv2.MORPH_OPEN, kernel)
+    
+    # 2. Fill small holes with closing (dilation followed by erosion)
+    cleaned = cv2.morphologyEx(cleaned, cv2.MORPH_CLOSE, kernel)
+    
+    # 3. Remove small connected components (islands) with better handling of thin roads
+    num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(cleaned, connectivity=CONNECTIVITY)
+    
+    # Create output mask
+    cleaned_mask = np.zeros_like(cleaned)
+    
+    # Keep components based on area AND shape characteristics
+    for i in range(1, num_labels):  # Start from 1 to skip background
+        area = stats[i, cv2.CC_STAT_AREA]
+        width = stats[i, cv2.CC_STAT_WIDTH]
+        height = stats[i, cv2.CC_STAT_HEIGHT]
+        
+        # Calculate aspect ratio (length/width)
+        aspect_ratio = max(width, height) / max(min(width, height), 1)
+        
+        # Calculate perimeter (approximate)
+        perimeter = 2 * (width + height)
+        
+        # Keep if:
+        # 1. Area is large enough, OR
+        # 2. It's a long, thin structure (high aspect ratio and reasonable perimeter)
+        if (area >= min_area or 
+            (aspect_ratio >= 3 and perimeter >= 20 and area >= 5)):
+            cleaned_mask[labels == i] = 1
+    
+    # 4. Apply additional dilation to connect nearby road segments
+    if kernel_size > 1:
+        cleaned_mask = cv2.dilate(cleaned_mask, kernel, iterations=1)
+        # Clean up again after dilation
+        cleaned_mask = cv2.morphologyEx(cleaned_mask, cv2.MORPH_CLOSE, kernel)
+    
+    return cleaned_mask.astype(np.uint8) * 255
+
+def postprocess(logits):
+    """
+    Postprocess the model output to create a clean road mask.
+    
+    Args:
+        logits: Model output tensor
+        
+    Returns:
+        dict: Contains the processed mask and metadata
+    """
+    # Convert logits to probabilities
+    if isinstance(logits, torch.Tensor):
+        # Apply sigmoid if not already applied
+        if logits.max() > 1.0:
+            probabilities = torch.sigmoid(logits)
+        else:
+            probabilities = logits
+        
+        # Convert to numpy
+        mask_np = probabilities.squeeze().cpu().numpy()
+    else:
+        mask_np = logits
+    
+    # Threshold the mask
+    binary_mask = (mask_np > THRESHOLD).astype(np.uint8) * 255
+    
+    # Clean the mask using morphological operations
+    cleaned_mask = _clean_road_mask(binary_mask)
+    
+    # Convert to PIL Image for easier handling
+    mask_image = Image.fromarray(cleaned_mask)
+    
+    # Calculate statistics
+    road_pixels = np.sum(cleaned_mask > 0)
+    total_pixels = cleaned_mask.size
+    road_percentage = (road_pixels / total_pixels) * 100
+    
+    # Create result dictionary
+    result = {
+        "mask": mask_image,
+        "road_percentage": round(road_percentage, 2),
+        "road_pixels": int(road_pixels),
+        "total_pixels": int(total_pixels),
+        "threshold_used": THRESHOLD,
+        "mask_shape": cleaned_mask.shape
+    }
+    
+    return result
+
+@torch.inference_mode()
+def predict(raw_input):
+    """
+    Main prediction function that processes input and returns road detection results.
+    
+    Args:
+        raw_input: Input image (base64, file path, or PIL Image)
+        
+    Returns:
+        dict: Road detection results with mask and metadata
+    """
+    try:
+        # Preprocess input
+        x = preprocess(raw_input)
+        
+        # Run model inference
+        logits = MODEL(x)
+        
+        # Postprocess results
+        result = postprocess(logits)
+        
+        return result
+        
+    except Exception as e:
+        return {
+            "error": str(e),
+            "status": "failed"
+        }
+
+def gradio_ui():
+    """Create Gradio interface for road detection."""
+    with gr.Blocks(title="Road Detection Model") as demo:
+        gr.Markdown("# 🛣️ Road Detection Model")
+        gr.Markdown("Upload a satellite image to detect roads.")
+        
+        with gr.Row():
+            with gr.Column():
+                # Input
+                input_image = gr.Image(
+                    label="Upload Satellite Image",
+                    type="pil",
+                    height=400
+                )
+                
+                # Parameters
+                with gr.Accordion("Advanced Parameters", open=False):
+                    threshold = gr.Slider(
+                        minimum=0.1,
+                        maximum=0.99,
+                        value=THRESHOLD,
+                        step=0.01,
+                        label="Detection Threshold"
+                    )
+                    min_area = gr.Slider(
+                        minimum=10,
+                        maximum=100,
+                        value=MIN_AREA,
+                        step=5,
+                        label="Minimum Road Area (pixels)"
+                    )
+                
+                # Run button
+                run_btn = gr.Button("🚀 Detect Roads", variant="primary")
+            
+            with gr.Column():
+                # Output
+                output_image = gr.Image(
+                    label="Detected Roads",
+                    height=400
+                )
+                
+                # Statistics
+                with gr.Accordion("Detection Statistics", open=True):
+                    road_percentage = gr.Number(
+                        label="Road Coverage (%)",
+                        precision=2
+                    )
+                    road_pixels = gr.Number(
+                        label="Road Pixels",
+                        precision=0
+                    )
+                    total_pixels = gr.Number(
+                        label="Total Pixels",
+                        precision=0
+                    )
+        
+        # Example images
+        gr.Examples(
+            examples=[
+                ["examples/example.jpg"]
+            ],
+            inputs=input_image,
+            label="Example Image"
+        )
+        
+        # Define prediction function
+        def predict_with_params(image, thresh, area):
+            global THRESHOLD, MIN_AREA
+            THRESHOLD = thresh
+            MIN_AREA = int(area)
+            
+            if image is None:
+                return None, 0, 0, 0
+            
+            result = predict(image)
+            
+            if "error" in result:
+                return None, 0, 0, 0
+            
+            return (
+                result["mask"],
+                result["road_percentage"],
+                result["road_pixels"],
+                result["total_pixels"]
+            )
+        
+        # Connect components
+        run_btn.click(
+            fn=predict_with_params,
+            inputs=[input_image, threshold, min_area],
+            outputs=[output_image, road_percentage, road_pixels, total_pixels]
+        )
+    
+    return demo
+
+if __name__ == "__main__":
+    load_model()
+    ui = gradio_ui()
+    ui.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))
+else:
+    # For Spaces
+    load_model()
+    demo = gradio_ui()

model_def.py

@@ -0,0 +1,61 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# --- UNet Model Definition ---
+class UNet(nn.Module):
+    def __init__(self, in_channels: int = 3, out_channels: int = 1, features: List[int] = [64, 128, 256, 512]):
+        super(UNet, self).__init__()
+        self.encoder = nn.ModuleList()
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.decoder = nn.ModuleList()
+
+        # Encoder
+        for feature in features:
+            self.encoder.append(self._conv_block(in_channels, feature))
+            in_channels = feature
+
+        # Bottleneck
+        self.bottleneck = self._conv_block(features[-1], features[-1]*2)
+
+        # Decoder
+        for feature in reversed(features):
+            self.decoder.append(
+                nn.ConvTranspose2d(feature*2, feature, kernel_size=2, stride=2)
+            )
+            self.decoder.append(self._conv_block(feature*2, feature))
+
+        self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)
+
+    def forward(self, x):
+        skip_connections = []
+        for down in self.encoder:
+            x = down(x)
+            skip_connections.append(x)
+            x = self.pool(x)
+        x = self.bottleneck(x)
+        skip_connections = skip_connections[::-1]
+        for idx in range(0, len(self.decoder), 2):
+            x = self.decoder[idx](x)
+            skip_connection = skip_connections[idx//2]
+            if x.shape != skip_connection.shape:
+                x = F.interpolate(x, size=skip_connection.shape[2:])
+            x = torch.cat((skip_connection, x), dim=1)
+            x = self.decoder[idx+1](x)
+        return torch.sigmoid(self.final_conv(x))
+
+    @staticmethod
+    def _conv_block(in_channels, out_channels):
+        return nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+        )
+
+
+def build_model():
+    # If you need custom args (e.g., from a config.json), read & pass them here.
+    return UNet()

push_to_hf.py

@@ -0,0 +1,54 @@
+from dotenv import load_dotenv
+
+load_dotenv()
+
+
+
+import os
+import shutil
+from huggingface_hub import HfApi, create_repo, upload_folder
+
+# ---- Setup ----
+# 1) Put your token in env:  export HUGGING_FACE_INFERENCE_TOKEN=hf_************************
+# 2) Choose target: either your username or an organization you belong to.
+HF_TOKEN = os.environ.get("HUGGING_FACE_INFERENCE_TOKEN")
+assert HF_TOKEN, "Set HUGGINGFACE_TOKEN env var."
+
+# Change these:
+SPACE_OWNER = "dunedain-ai"          
+SPACE_NAME  = "road-detection-model"             
+SPACE_SDK   = "gradio"                        # Space SDK: gradio | streamlit | static | ...
+REPO_ID     = f"{SPACE_OWNER}/{SPACE_NAME}"
+
+# Folder containing your app.py etc.
+LOCAL_DIR = os.path.dirname(os.path.abspath(__file__))
+
+def main():
+    print('Accessing Hugging Face...')
+    api = HfApi(token=HF_TOKEN)
+    print('Hugging Face Accessed!')
+
+    # Create Space repo (idempotent; set exist_ok=True)
+    print('Creating Repo...')
+    create_repo(
+        repo_id=REPO_ID,
+        repo_type="space",
+        space_sdk=SPACE_SDK,
+        private=False,     # set True for private
+        exist_ok=True
+    )
+    print('Repo Created!')
+
+
+    # Upload everything in this directory
+    upload_folder(
+        repo_id=REPO_ID,
+        repo_type="space",
+        folder_path=LOCAL_DIR,
+        commit_message="Initial Space upload"
+    )
+
+    print(f"✅ Uploaded. Space: https://huggingface.co/spaces/{REPO_ID}")
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,33 @@
+# Core ML and Deep Learning
+torch>=2.0.0
+torchvision>=0.15.0
+numpy>=1.21.0
+
+# Computer Vision
+opencv-python>=4.8.0
+Pillow>=9.0.0
+
+# Web Interface
+gradio>=4.0.0
+
+# Image Processing
+scikit-image>=0.20.0
+
+# Utilities
+requests>=2.28.0
+python-dotenv>=1.0.0
+
+# Optional: For Hugging Face integration
+transformers>=4.30.0
+huggingface-hub>=0.16.0
+
+# Optional: For advanced image processing
+scipy>=1.10.0
+
+# Optional: For better performance
+accelerate>=0.20.0
+
+# Development and testing (optional)
+pytest>=7.0.0
+black>=23.0.0
+flake8>=6.0.0

runtime.txt

@@ -0,0 +1 @@
+python-3.12

weights/road_detection_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4f2ccb110c1c5b7eaca81b1ee136737cc8c1a0feccde2b0b9721a9a7fd5a09e5
+size 124236775