test handler

handler.py

@@ -0,0 +1,49 @@
+from typing import Any, Dict, List
+import base64
+import io
+import tempfile
+from PIL import Image
+from pdf2image import convert_from_path
+
+
+def process_image(pdfdata):
+    prompt = "detect the accident diagram"
+    # PDF input -> convert to images -> find crash diagram
+    with tempfile.NamedTemporaryFile(delete=False) as temp_file:
+        temp_file.write(pdfdata)
+        temp_file.flush()
+        print("temporary name:", temp_file.name)
+        images = convert_from_path(temp_file.name)
+
+    crash_img = None
+    for img in images:
+        crash_img = img.crop((0,0, 200, 200))
+        break
+
+    if crash_img:
+        img_io = io.BytesIO()
+        crash_img.save(img_io, "PNG")
+        img_io.seek(0)
+        # Return the cropped image as base64
+        return {"data": base64.b64encode(img_io.read()).decode("utf-8"), "mime_type": "image/png"}
+    else:
+        return {"error": "No crash diagram detected"}
+
+
+class EndpointHandler:
+     def __init__(self, path: str = ""):
+        """Initialize the endpoint handler.
+        
+        Args:
+            path: Path to the model artifacts
+        """
+        print("initialized")
+        pass
+     
+     def __call__(self, data: Any) -> List[List[Dict[str, str]]]:
+         pdfdata = data.pop("filedata", data)
+         if isinstance(pdfdata, str):
+             print("decoding pdfdata")
+             pdfdata = base64.b64decode(pdfdata)
+
+         return process_image(pdfdata)

requirements.txt

@@ -0,0 +1,2 @@
+pdf2image==1.17.0
+pillow==10.4.0