sdfgs

handler.py

@@ -1,9 +1,19 @@
 from typing import  Dict, List, Any
+import base64
+import io
 import torch
 import numpy as np
 import torch.nn.functional as F
 from serkan import SimpleUpscaleModel
 import os
+from PIL import Image
+
+def decode_image(self, base64_str: str) -> np.ndarray:
+        """Decode base64 string to an image (numpy array)"""
+        image_data = base64.b64decode(base64_str)
+        image = Image.open(io.BytesIO(image_data))
+        return np.array(image)
+
 class EndpointHandler():
     def __init__(self, path="."):
         # load the optimized model
@@ -12,7 +22,7 @@ class EndpointHandler():
         self.model.load_state_dict(torch.load(model_path))
 
 
-
+    
 
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         """
@@ -26,10 +36,9 @@ class EndpointHandler():
         """
         inputs = data.pop("inputs", data)
         img = inputs["image"]
-
+        img = self.decode_image(img)
+        img = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0).float()
         # Load the image
-        img = np.float32(img)
-
         upscaled = self.model(img)
         # postprocess the prediction
         return "OKAY"