Update handler.py

handler.py

@@ -5,6 +5,7 @@ from PIL import Image
 from huggingface_hub import hf_hub_download
 import io
 import base64
+import numpy as np
 
 # --- Basic UNet Components ---
 class DoubleConv(nn.Module):
@@ -51,13 +52,13 @@ class Up(nn.Module):
 
 class OutConv(nn.Module):
     def __init__(self, in_channels, out_channels):
-        super(OutConv, self).__init__()
+        super().__init__()
         self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
 
     def forward(self, x):
         return self.conv(x)
 
-# --- UNet Architecture ---
+# --- Full UNet ---
 class UNet(nn.Module):
     def __init__(self, n_channels=3, n_classes=1, bilinear=True):
         super().__init__()
@@ -90,12 +91,13 @@ class UNet(nn.Module):
         logits = self.outc(x)
         return torch.sigmoid(logits)
 
-# --- Endpoint Handler ---
+# --- EndpointHandler for Hugging Face Inference Endpoint ---
 class EndpointHandler:
     def __init__(self, path=""):
         model_path = hf_hub_download(repo_id="whitney0507/unet-model", filename="UNet_Model.pth")
         self.model = UNet()
-        self.model.load_state_dict(torch.load(model_path, map_location="cpu"))
+        state_dict = torch.load(model_path, map_location=torch.device("cpu"))
+        self.model.load_state_dict(state_dict)
         self.model.eval()
         self.transform = transforms.Compose([
             transforms.Resize((256, 256)),
@@ -109,9 +111,11 @@ class EndpointHandler:
 
         with torch.no_grad():
             output = self.model(input_tensor)
-            mask = (output > 0.5).int().squeeze().byte().cpu().numpy()
+            mask = (output > 0.5).int().squeeze().cpu().numpy()
 
-        result_img = Image.fromarray(mask * 255)
+        # Ensure mask is in uint8 format for image encoding
+        result_img = Image.fromarray((mask * 255).astype(np.uint8))
         buffer = io.BytesIO()
         result_img.save(buffer, format="PNG")
-        return {"prediction": base64.b64encode(buffer.getvalue()).decode("utf-8")}
+        encoded_output = base64.b64encode(buffer.getvalue()).decode("utf-8")
+        return {"prediction": encoded_output}