Upload handler.py (#4)

- Upload handler.py (4f6f56e6ade5adb4fc5e89079d346f6007fffa0c)

handler.py

@@ -3,12 +3,15 @@ from transformers import BlipProcessor, BlipForConditionalGeneration
 from PIL import Image
 from io import BytesIO
 import torch
+import base64
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
+
 class EndpointHandler():
     def __init__(self, path=""):
-        self.model = BlipForConditionalGeneration.from_pretrained("quadranttechnologies/qhub-blip-image-captioning-finetuned").to(device)
+        self.model = BlipForConditionalGeneration.from_pretrained(
+            "quadranttechnologies/qhub-blip-image-captioning-finetuned").to(device)
         self.processor = BlipProcessor.from_pretrained("quadranttechnologies/qhub-blip-image-captioning-finetuned")
         self.model.eval()
         self.model = self.model.to(device).to(device)
@@ -27,9 +30,9 @@ class EndpointHandler():
         text = data.get("text", "")
         parameters = data.pop("parameters", {})
 
-        raw_images = Image.open(BytesIO(inputs)).convert("")
+        raw_images = Image.open(BytesIO(base64.b64decode(inputs))).convert("RGB")
 
-        processed_image = self.processor(images=raw_images, text = text, return_tensors="pt")
+        processed_image = self.processor(images=raw_images, text=text, return_tensors="pt")
         processed_image["pixel_values"] = processed_image["pixel_values"].to(device)
         processed_image = {**processed_image, **parameters}
 
@@ -41,4 +44,5 @@ class EndpointHandler():
 
         return {"description": description}
 
+
 handler = EndpointHandler()