GMI-AI/Road-segmentation-app

1"""2Interactive Gradio app for road segmentation using U-Net model.3Upload aerial/satellite images and get road segmentation predictions.4"""5 6import gradio as gr7import torch8import torch.nn as nn9import numpy as np10from PIL import Image11import torchvision.transforms as transforms12import os13 14 15# MODEL DEFINITION (U-Net Architecture)16 17 18class DoubleConv(nn.Module):19    def __init__(self, in_channels, out_channels, mid_channels=None):20        super().__init__()21        if not mid_channels:22            mid_channels = out_channels23        self.double_conv = nn.Sequential(24            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),25            nn.BatchNorm2d(mid_channels),26            nn.ReLU(inplace=True),27            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),28            nn.BatchNorm2d(out_channels),29            nn.ReLU(inplace=True)30        )31 32    def forward(self, x):33        return self.double_conv(x)34 35 36class Down(nn.Module):37    def __init__(self, in_channels, out_channels):38        super().__init__()39        self.maxpool_conv = nn.Sequential(40            nn.MaxPool2d(2),41            DoubleConv(in_channels, out_channels)42        )43 44    def forward(self, x):45        return self.maxpool_conv(x)46 47 48class Up(nn.Module):49    def __init__(self, in_channels, out_channels, bilinear=True):50        super().__init__()51        if bilinear:52            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)53            self.conv = DoubleConv(in_channels, out_channels, in_channels // 2)54        else:55            self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)56            self.conv = DoubleConv(in_channels, out_channels)57 58    def forward(self, x1, x2):59        x1 = self.up(x1)60        diffY = x2.size()[2] - x1.size()[2]61        diffX = x2.size()[3] - x1.size()[3]62        x1 = nn.functional.pad(x1, [diffX // 2, diffX - diffX // 2,63                                     diffY // 2, diffY - diffY // 2])64        x = torch.cat([x2, x1], dim=1)65        return self.conv(x)66 67 68class OutConv(nn.Module):69    def __init__(self, in_channels, out_channels):70        super(OutConv, self).__init__()71        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)72 73    def forward(self, x):74        return self.conv(x)75 76 77class UNet(nn.Module):78    def __init__(self, n_channels=3, n_classes=2, bilinear=False):79        super(UNet, self).__init__()80        self.n_channels = n_channels81        self.n_classes = n_classes82        self.bilinear = bilinear83 84        self.inc = DoubleConv(n_channels, 64)85        self.down1 = Down(64, 128)86        self.down2 = Down(128, 256)87        self.down3 = Down(256, 512)88        factor = 2 if bilinear else 189        self.down4 = Down(512, 1024 // factor)90        self.up1 = Up(1024, 512 // factor, bilinear)91        self.up2 = Up(512, 256 // factor, bilinear)92        self.up3 = Up(256, 128 // factor, bilinear)93        self.up4 = Up(128, 64, bilinear)94        self.outc = OutConv(64, n_classes)95 96    def forward(self, x):97        x1 = self.inc(x)98        x2 = self.down1(x1)99        x3 = self.down2(x2)100        x4 = self.down3(x3)101        x5 = self.down4(x4)102        x = self.up1(x5, x4)103        x = self.up2(x, x3)104        x = self.up3(x, x2)105        x = self.up4(x, x1)106        logits = self.outc(x)107        return logits108 109 110 111# LOAD MODEL112 113 114device = torch.device('cpu')115 116# Initialize model117model = UNet(n_channels=3, n_classes=2, bilinear=True)118 119# Load weights120checkpoint_path = 'unet_continued_best.pth'121model.load_state_dict(torch.load(checkpoint_path, map_location=device))122model.to(device)123model.eval()124 125print(f"✓ Model loaded successfully on {device}")126 127 128 129# PREDICTION FUNCTION130 131 132def predict_road(image):133    """134    Predict road segmentation from input image.135    136    Args:137        image: PIL Image or numpy array138        139    Returns:140        Original image, Road mask, Overlay visualization141    """142    143    # Convert to PIL if needed144    if isinstance(image, np.ndarray):145        image = Image.fromarray(image)146    147    # Resize to model input size148    original_size = image.size149    image_resized = image.resize((256, 256), Image.BILINEAR)150    151    # Prepare image for model152    transform = transforms.Compose([153        transforms.ToTensor(),154        transforms.Normalize(mean=[0.485, 0.456, 0.406], 155                           std=[0.229, 0.224, 0.225])156    ])157    158    input_tensor = transform(image_resized).unsqueeze(0).to(device)159    160    # Predict161    with torch.no_grad():162        output = model(input_tensor)163        prediction = torch.argmax(output, dim=1).squeeze().cpu().numpy()164    165    # Resize prediction back to original size166    prediction_resized = Image.fromarray((prediction * 255).astype(np.uint8))167    prediction_resized = prediction_resized.resize(original_size, Image.NEAREST)168    prediction_resized = np.array(prediction_resized) / 255.0169    170    # Create visualizations171    original_np = np.array(image)172    173    # Road mask (binary)174    road_mask = (prediction_resized * 255).astype(np.uint8)175    road_mask_colored = np.zeros((*road_mask.shape, 3), dtype=np.uint8)176    road_mask_colored[road_mask > 0] = [255, 0, 0]  # Red for roads177    178    # Overlay179    overlay = original_np.copy()180    overlay[prediction_resized > 0.5] = overlay[prediction_resized > 0.5] * 0.5 + np.array([255, 0, 0]) * 0.5181    overlay = overlay.astype(np.uint8)182    183    return original_np, road_mask_colored, overlay184 185 186 187# GRADIO INTERFACE188 189 190# Examples191examples = [192    ["examples/example1.tiff"] ,193    ["examples/example2.tiff"] ,194]195examples = [ex for ex in examples if ex is not None]196 197# Create interface198demo = gr.Interface(199    fn=predict_road,200    inputs=gr.Image(type="pil", label="Upload Aerial/Satellite Image"),201    outputs=[202        gr.Image(type="numpy", label="Original Image"),203        gr.Image(type="numpy", label="Road Segmentation (Red)"),204        gr.Image(type="numpy", label="Overlay")205    ],206    title="🛣️ Road Segmentation from Aerial Images",207    description="""208    Upload an aerial or satellite image to detect roads using a U-Net deep learning model.209    210    **How to use:**211    1. Click **Clear** to remove the displayed example.212    2. Upload an satellite image 213    3. Click **Submit**214    4. View the detected roads highlighted in red215    216    **Model:** U-Net trained on Massachusetts Roads Dataset217    **Performance:** IoU ~0.55, F1 ~0.70218 219    *Note: Examples (images) are available below.220    """,221    examples=examples if examples else None,222)223 224if __name__ == "__main__":225    demo.launch()