File size: 1,432 Bytes
f3abe3d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 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) |