pdf

handler.py

@@ -4,6 +4,8 @@ import torch
 from PIL import Image
 import io
 import base64
+import fitz  # PyMuPDF
+import tempfile
 
 class EndpointHandler():
     def __init__(self, path=""):
@@ -19,20 +21,8 @@ class EndpointHandler():
         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")
-        
+    def process_image(self, image):
+        """Process a single image and return extractions"""
         encoding = self.processor(
             image,
             truncation=True,
@@ -54,7 +44,85 @@ class EndpointHandler():
             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]}
+                    "bbox": {
+                        "x": box[0], 
+                        "y": box[1], 
+                        "width": box[2] - box[0], 
+                        "height": box[3] - box[1]
+                    }
                 })
         
-        return {"extractions": results}
+        return results
+    
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        inputs = data.pop("inputs", data)
+        
+        # Handle different input formats
+        if isinstance(inputs, dict):
+            # Check if it's a PDF
+            if "pdf" in inputs:
+                file_data = inputs["pdf"]
+            else:
+                file_data = inputs.get("image", inputs.get("inputs", ""))
+        else:
+            file_data = inputs
+        
+        # Remove base64 prefix if present
+        if isinstance(file_data, str) and "base64," in file_data:
+            file_data = file_data.split("base64,")[1]
+        
+        # Decode base64
+        file_bytes = base64.b64decode(file_data)
+        
+        # Check if it's a PDF or image
+        if file_bytes.startswith(b'%PDF'):
+            # Process PDF
+            all_results = []
+            
+            # Save PDF to temporary file
+            with tempfile.NamedTemporaryFile(suffix='.pdf', delete=False) as tmp_file:
+                tmp_file.write(file_bytes)
+                tmp_file.flush()
+                
+                # Open with PyMuPDF
+                pdf_document = fitz.open(tmp_file.name)
+                
+                # Process each page
+                for page_num in range(len(pdf_document)):
+                    page = pdf_document[page_num]
+                    
+                    # Convert page to image (PIL format)
+                    mat = fitz.Matrix(2.0, 2.0)  # 2x scaling for better quality
+                    pix = page.get_pixmap(matrix=mat)
+                    img_data = pix.tobytes("png")
+                    image = Image.open(io.BytesIO(img_data)).convert("RGB")
+                    
+                    # Process the page
+                    page_results = self.process_image(image)
+                    
+                    # Add page context to results
+                    all_results.append({
+                        "page": page_num + 1,
+                        "page_width": page.rect.width,
+                        "page_height": page.rect.height,
+                        "extractions": page_results
+                    })
+                
+                pdf_document.close()
+            
+            # Return all pages' results
+            return {
+                "document_type": "pdf",
+                "total_pages": len(all_results),
+                "pages": all_results
+            }
+        
+        else:
+            # Process as image
+            image = Image.open(io.BytesIO(file_bytes)).convert("RGB")
+            results = self.process_image(image)
+            
+            return {
+                "document_type": "image",
+                "extractions": results
+            }

requirements.txt

@@ -2,3 +2,4 @@ transformers>=4.35.0
 torch>=2.0.0
 pillow>=9.0.0
 pytesseract>=0.3.10
+PyMuPDF>=1.23.0