update handler

Dockerfile

@@ -0,0 +1,29 @@
+FROM python:3.9-slim-bullseye AS base
+WORKDIR /app
+
+FROM base AS build
+COPY ["requirements.txt", "./"]
+RUN \
+  apt-get update && \
+  apt-get install -y --no-install-recommends \
+    git \
+    g++ && \
+  rm -rf /var/lib/apt/lists/* && \
+  # Install dependencies.
+  python -m venv ./venv && \
+  ./venv/bin/pip install --upgrade \
+    pip \
+    setuptools \
+    wheel && \
+  ./venv/bin/pip install --no-cache-dir -r ./requirements.txt && \
+  ./venv/bin/pip install --no-cache-dir 'git+https://github.com/facebookresearch/detectron2.git@d1e04565d3bec8719335b88be9e9b961bf3ec464'
+
+FROM base AS final
+RUN \
+  apt-get update && \
+  apt-get install -y --no-install-recommends \
+    tesseract-ocr && \
+  rm -rf /var/lib/apt/lists/*
+COPY --from=build ["/app/venv", "./venv"]
+# Copy the source code in last to optimize rebuilding the image.
+COPY [".", "./"]

handler.py

@@ -21,8 +21,8 @@ class EndpointHandler:
         self.PROCESSOR = LayoutLMv3Processor(self.FEATURE_EXTRACTOR, self.TOKENIZER)
         self.MODEL = LayoutLMv3ForSequenceClassification.from_pretrained("OtraBoi/document_classifier_testing").to(device)
 
-    def __call__(self, data: bytes):
-        image = Image.open(io.BytesIO(data)).convert("RGB")
+    def __call__(self, data: Dict):
+        image = Image.open(io.BytesIO(data["inputs"])).convert("RGB")
         encoding = self.PROCESSOR(image, return_tensors="pt", padding="max_length", truncation=True)
         
         for k,v in encoding.items():