Create app.py

app.py

@@ -0,0 +1,42 @@
+import torch
+import gradio as gr
+from torchvision import transforms
+from PIL import Image
+from huggingface_hub import hf_hub_download
+
+# 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)
+
+model_path = hf_hub_download(repo_id="TheArchitect416/oil-spill-segmentation-model", filename="model.pth")
+model = MySegmentationModel()
+model.load_state_dict(torch.load(model_path, map_location="cpu"))
+model.eval()
+
+# Define preprocessing
+transform = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.ToTensor()
+])
+
+# Define inference function
+def predict(image):
+    image = transform(image).unsqueeze(0)  # Add batch dimension
+    with torch.no_grad():
+        output = model(image)
+    return output.squeeze(0).numpy()  # Convert to numpy for visualization
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs="image"
+)
+
+iface.launch()