Update ai_mapping.py

ai_mapping.py

@@ -1,7 +1,7 @@
 from transformers import LayoutLMv3Tokenizer, LayoutLMv3ForTokenClassification, LayoutLMv3ImageProcessor
 import torch
 from PIL import Image
-import pdf2image
+import fitz  # PyMuPDF
 from typing import Dict, List
 import os
 from huggingface_hub import login
@@ -33,8 +33,14 @@ def extract_key_values_with_layoutlm(text_data: str, pdf_path: str) -> Dict[str,
             key_values.update({"Date": dates[0] if dates else "", "Amount": amounts[0] if amounts else ""})
 
         # Attempt LayoutLMv3 processing
-        images = pdf2image.convert_from_path(pdf_path)
-        for i, image in enumerate(images):
+        doc = fitz.open(pdf_path)
+        for page_num in range(len(doc)):
+            page = doc[page_num]
+            pix = page.get_pixmap(matrix=fitz.Matrix(300/72, 300/72))  # 300 DPI
+            img_path = f"{pdf_path}_page_{page_num}.png"
+            pix.save(img_path)
+            image = Image.open(img_path)
+
             encoding = feature_extractor(images=[image], text=text_data.splitlines(), return_tensors="pt")
             input_ids = encoding["input_ids"]
             attention_mask = encoding["attention_mask"]
@@ -58,6 +64,11 @@ def extract_key_values_with_layoutlm(text_data: str, pdf_path: str) -> Dict[str,
             if current_key and current_value:
                 key_values[current_key] = " ".join(current_value).strip()
 
+            # Clean up temporary image
+            if os.path.exists(img_path):
+                os.unlink(img_path)
+
+        doc.close()
         return key_values if key_values else {"status": "failed", "error": "No key-value pairs extracted", "key_values": {}}
     except Exception as e:
         return {"status": "failed", "error": str(e), "key_values": {}}