Create handler.py

handler.py

@@ -0,0 +1,88 @@
+import torch
+import segmentation_models_pytorch as smp
+from torchvision import transforms
+from PIL import Image
+import io
+import json
+import base64
+import numpy as np
+
+# Define the number of output classes (update if needed)
+NUM_CLASSES = 4
+
+# Define preprocessing transforms (should match what was used during training)
+preprocess = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=(0.485, 0.456, 0.406),  # ImageNet means
+                         std=(0.229, 0.224, 0.225))
+])
+
+# Define class-color mapping for segmentation mask visualization
+COLOR_MAPPING = {
+    0: [0, 0, 0],        # Background
+    1: [255, 0, 124],    # Oil
+    2: [255, 204, 51],   # Others
+    3: [51, 221, 255]    # Water
+}
+
+def colorize_mask(mask):
+    """Convert a 2D segmentation mask into an RGB image."""
+    h, w = mask.shape
+    color_mask = np.zeros((h, w, 3), dtype=np.uint8)
+    for cls, color in COLOR_MAPPING.items():
+        color_mask[mask == cls] = color
+    return color_mask
+
+class OilSpillSegmentationHandler:
+    def __init__(self):
+        """Load the model and set it to evaluation mode."""
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = smp.Unet(
+            encoder_name="resnet34",  # Ensure this matches your training
+            encoder_weights=None,     # Weights are loaded from state_dict
+            in_channels=3,
+            classes=NUM_CLASSES
+        )
+        self.model.load_state_dict(torch.load("model.pth", map_location=self.device))
+        self.model.to(self.device)
+        self.model.eval()
+
+    def preprocess(self, image_bytes):
+        """Preprocess input image (convert to tensor)."""
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        image_tensor = preprocess(image).unsqueeze(0).to(self.device)
+        return image_tensor, image
+
+    def inference(self, image_tensor):
+        """Run inference and return the segmentation mask."""
+        with torch.no_grad():
+            output = self.model(image_tensor)
+            pred_mask = torch.argmax(output, dim=1).squeeze(0).cpu().numpy().astype(np.uint8)
+        return pred_mask
+
+    def postprocess(self, pred_mask):
+        """Convert segmentation mask to colorized image."""
+        colorized_mask = colorize_mask(pred_mask)
+        return Image.fromarray(colorized_mask)
+
+    def handle_request(self, request_body):
+        """Handle API request: preprocess, infer, postprocess."""
+        try:
+            data = json.loads(request_body)
+            image_bytes = base64.b64decode(data["image"])
+            image_tensor, original_image = self.preprocess(image_bytes)
+            pred_mask = self.inference(image_tensor)
+            output_image = self.postprocess(pred_mask)
+
+            # Convert output image to base64
+            buffered = io.BytesIO()
+            output_image.save(buffered, format="PNG")
+            output_b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+            return json.dumps({"output_image": output_b64})
+        except Exception as e:
+            return json.dumps({"error": str(e)})
+
+# Instantiate the handler
+handler = OilSpillSegmentationHandler()