|
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" |
|
|
|
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 {"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) |
