sfgsdfgs

handler.py

@@ -40,5 +40,19 @@ class EndpointHandler():
         img = torch.from_numpy(img).permute(2, 0, 1).unsqueeze(0).float()
         # Load the image
         upscaled = self.model(img)
+        upscaled = upscaled.squeeze(0).permute(1,2,0)
+        upscaled = upscaled.numpy()
+        
+        pil = Image.fromarray(upscaled)
+        # Save the image to a buffer
+        buffered = io.BytesIO()
+        pil.save(buffered, format="PNG")
+        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+        # Return a dictionary with the base64 image and additional data
+        return {
+            "image": img_str
+        }
+
         # postprocess the prediction
         return "OKAY"