imbkarthi/cloud-segmentation

1import gradio as gr2import torch3import torch.nn as nn4import torch.nn.functional as F5import numpy as np6from PIL import Image7import torchvision.transforms as T8 9# --- U-Net Architecture ---10class DoubleConv(nn.Module):11    def __init__(self, in_channels, out_channels):12        super().__init__()13        self.conv = nn.Sequential(14            nn.Conv2d(in_channels, out_channels, 3, padding=1),15            nn.BatchNorm2d(out_channels),16            nn.ReLU(inplace=True),17            nn.Conv2d(out_channels, out_channels, 3, padding=1),18            nn.BatchNorm2d(out_channels),19            nn.ReLU(inplace=True)20        )21    def forward(self, x): return self.conv(x)22 23class UNet(nn.Module):24    def __init__(self, n_channels=3, n_classes=4):25        super(UNet, self).__init__()26        self.inc = DoubleConv(n_channels, 64)27        self.down1 = nn.Sequential(nn.MaxPool2d(2), DoubleConv(64, 128))28        self.down2 = nn.Sequential(nn.MaxPool2d(2), DoubleConv(128, 256))29        self.up1 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)30        self.conv_up1 = DoubleConv(256, 128)31        self.up2 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)32        self.conv_up2 = DoubleConv(128, 64)33        self.outc = nn.Conv2d(64, n_classes, kernel_size=1)34 35    def forward(self, x):36        x1 = self.inc(x)37        x2 = self.down1(x1)38        x3 = self.down2(x2)39        x = self.up1(x3)40        x = torch.cat([x, x2], dim=1)41        x = self.conv_up1(x)42        x = self.up2(x)43        x = torch.cat([x, x1], dim=1)44        x = self.conv_up2(x)45        return torch.sigmoid(self.outc(x))46 47# --- Model & Logic ---48LABELS = ["Sugar", "Flower", "Fish", "Gravel"]49COLORS = {"Sugar": "#FFFFFF", "Flower": "#FFD700", "Fish": "#1E90FF", "Gravel": "#8B4513"}50 51device = "cuda" if torch.cuda.is_available() else "cpu"52model = UNet(n_classes=4).to(device)53model.eval()54 55def predict(input_img):56    if input_img is None: return None57    58    # Preprocess59    h, w = input_img.shape[:2]60    transform = T.Compose([T.ToPILImage(), T.Resize((256, 256)), T.ToTensor()])61    input_tensor = transform(input_img).unsqueeze(0).to(device)62    63    with torch.no_grad():64        output = model(input_tensor)65    66    # Resize mask back to original image dimensions67    output = F.interpolate(output, size=(h, w), mode='bilinear').squeeze(0).cpu().numpy()68    69    annotations = []70    for i, label in enumerate(LABELS):71        # Using 0.5 threshold for the prototype72        mask = (output[i] > 0.5).astype(np.uint8)73        if mask.sum() > 0:74            annotations.append((mask, label))75            76    return (input_img, annotations)77 78# --- Gradio UI ---79demo = gr.Interface(80    fn=predict,81    inputs=gr.Image(label="Upload Satellite Image"),82    outputs=gr.AnnotatedImage(label="Cloud Segmentations", color_map=COLORS),83    title="Cloud Pattern Detector",84    description="Upload a satellite image to identify **Sugar, Flower, Fish, or Gravel** patterns."85)86 87if __name__ == "__main__":88    demo.launch()