Add custom handler for LayoutLMv3 inference

handler.py

@@ -0,0 +1,60 @@
+from typing import Dict, List, Any
+from transformers import LayoutLMv3Processor, LayoutLMv3ForTokenClassification
+import torch
+from PIL import Image
+import io
+import base64
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        # Load from Microsoft's repo
+        self.processor = LayoutLMv3Processor.from_pretrained(
+            "microsoft/layoutlmv3-base",
+            apply_ocr=True
+        )
+        self.model = LayoutLMv3ForTokenClassification.from_pretrained(
+            "microsoft/layoutlmv3-base"
+        )
+        self.model.eval()
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(self.device)
+    
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        inputs = data.pop("inputs", data)
+        
+        if isinstance(inputs, dict):
+            image_data = inputs.get("image", inputs.get("inputs", ""))
+        else:
+            image_data = inputs
+        
+        if "base64," in image_data:
+            image_data = image_data.split("base64,")[1]
+        
+        image_bytes = base64.b64decode(image_data)
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        
+        encoding = self.processor(
+            image,
+            truncation=True,
+            padding="max_length",
+            max_length=512,
+            return_tensors="pt"
+        )
+        
+        encoding = {k: v.to(self.device) for k, v in encoding.items() if isinstance(v, torch.Tensor)}
+        
+        with torch.no_grad():
+            outputs = self.model(**encoding)
+        
+        tokens = self.processor.tokenizer.convert_ids_to_tokens(encoding["input_ids"][0].cpu())
+        boxes = encoding["bbox"][0].cpu().tolist()
+        
+        results = []
+        for token, box in zip(tokens, boxes):
+            if token not in ['[CLS]', '[SEP]', '[PAD]']:
+                results.append({
+                    "text": token,
+                    "bbox": {"x": box[0], "y": box[1], "width": box[2] - box[0], "height": box[3] - box[1]}
+                })
+        
+        return {"extractions": results}

requirements.txt

@@ -0,0 +1,4 @@
+transformers>=4.35.0
+torch>=2.0.0
+pillow>=9.0.0
+pytesseract>=0.3.10