Update handler.py

handler.py

@@ -1,26 +1,50 @@
-from typing import Dict, List, Any
-from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+import base64
+import io
+from typing import Any, Dict, List
+
+import requests
 import torch
+from PIL import Image
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-class EndpointHandler():
+
+class EndpointHandler:
     def __init__(self, path=""):
-      self.processor = TrOCRProcessor.from_pretrained(path)
-      self.model = VisionEncoderDecoderModel.from_pretrained(path)
+        self.processor = TrOCRProcessor.from_pretrained(path)
+        self.model = VisionEncoderDecoderModel.from_pretrained(path)
+
+        self.model.to(device)
 
-      device = "cuda" if torch.cuda.is_available() else "cpu"
-      self.model.to(device)
-    
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         inputs = data.pop("inputs", data)
-        image_input = inputs.get('image')
+        image_input = inputs.get("image")
 
+        if not image_input:
+            return {"error": "No image provided."}
 
-        # process image
-        pixel_values = self.processor(images=image_input, return_tensors="pt").pixel_values
+        try:
+            if image_input.startswith("http"):
+                response = requests.get(image_input, stream=True)
+                if response.status_code == 200:
+                    image = Image.open(response.raw).convert("RGB")
+                else:
+                    return {
+                        "error": f"Failed to fetch image. Status code: {response.status_code}"
+                    }
+            else:
+                image_data = base64.b64decode(image_input)
+                image = Image.open(io.BytesIO(image_data)).convert("RGB")
+        except Exception as e:
+            return {"error": f"Failed to process the image. Details: {str(e)}"}
+
+        pixel_values = self.processor(images=image, return_tensors="pt").pixel_values
 
-        # run prediction
         generated_ids = self.model.generate(pixel_values.to(device))
-        
-        # decode output
-        prediction = generated_text = self.processor.batch_decode(generated_ids, skip_special_tokens=True)
-        return {"text":prediction[0]}
+
+        prediction = self.processor.batch_decode(
+            generated_ids, skip_special_tokens=True
+        )
+        return {"text": prediction[0]}