Update app.py

app.py

@@ -1,42 +1,60 @@
 import torch
 import gradio as gr
 from torchvision import transforms
-from PIL import Image
 from huggingface_hub import hf_hub_download
+import segmentation_models_pytorch as smp
+import numpy as np
 
-# Load model
-class MySegmentationModel(torch.nn.Module):
-    def __init__(self):
-        super(MySegmentationModel, self).__init__()
-        # Define your model architecture here
-        self.dummy_layer = torch.nn.Conv2d(3, 1, kernel_size=3, stride=1, padding=1)
-
-    def forward(self, x):
-        return self.dummy_layer(x)
+# Set the number of output classes (from your label_colors.txt, you have 4 classes)
+NUM_CLASSES = 4
 
+# Download the model state dictionary from your Hugging Face repository
 model_path = hf_hub_download(repo_id="TheArchitect416/oil-spill-segmentation-model", filename="model.pth")
-model = MySegmentationModel()
+
+# Create the model using segmentation_models_pytorch.
+# This should match the architecture you used during training.
+model = smp.Unet(
+    encoder_name="resnet34",       # for example, resnet34 was used in training
+    encoder_weights="imagenet",    # or you might have used pretrained weights from ImageNet
+    in_channels=3,                 # RGB images
+    classes=NUM_CLASSES            # number of segmentation classes
+)
+# Load the state dict (mapping the keys appropriately)
 model.load_state_dict(torch.load(model_path, map_location="cpu"))
 model.eval()
 
-# Define preprocessing
-transform = transforms.Compose([
+# Define preprocessing transforms (should match what was used during training)
+preprocess = transforms.Compose([
     transforms.Resize((256, 256)),
-    transforms.ToTensor()
+    transforms.ToTensor(),
+    transforms.Normalize(mean=(0.485, 0.456, 0.406),  # ImageNet means
+                         std=(0.229, 0.224, 0.225))
 ])
 
-# Define inference function
+# Define the inference function
 def predict(image):
-    image = transform(image).unsqueeze(0)  # Add batch dimension
+    """
+    Accepts a PIL image, preprocesses it, runs the model,
+    and returns the predicted mask.
+    """
+    # Preprocess the image
+    input_tensor = preprocess(image).unsqueeze(0)  # add batch dimension; shape: [1, 3, 256, 256]
+    
     with torch.no_grad():
-        output = model(image)
-    return output.squeeze(0).numpy()  # Convert to numpy for visualization
+        output = model(input_tensor)
+    
+    # The output is typically raw logits.
+    # Take argmax along the channel dimension to get the predicted class per pixel.
+    pred_mask = torch.argmax(output, dim=1).squeeze(0).cpu().numpy().astype(np.uint8)
+    return pred_mask
 
-# Create Gradio interface
+# Create a Gradio interface
 iface = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
-    outputs="image"
+    outputs=gr.Image(type="numpy"),
+    title="Oil Spill Segmentation",
+    description="Segment oil spills in aerial images."
 )
 
 iface.launch()