Update ui/ui_core.py

ui/ui_core.py

@@ -6,6 +6,10 @@ import json
 import gradio as gr
 from typing import List
 
+from transformers import LayoutLMv3Processor, LayoutLMv3ForTokenClassification
+from PIL import Image
+import torch
+
 # ✅ Fix: Add src to Python path
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "src")))
 
@@ -32,6 +36,22 @@ def clean_final_response(text: str) -> str:
         )
     return "".join(panels)
 
+def use_layoutlmv3_on_image(image_path):
+    processor = LayoutLMv3Processor.from_pretrained("microsoft/layoutlmv3-large")
+    model = LayoutLMv3ForTokenClassification.from_pretrained("microsoft/layoutlmv3-large")
+
+    image = Image.open(image_path).convert("RGB")
+    encoding = processor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**encoding)
+
+    logits = outputs.logits
+    predicted_class = logits.argmax(-1)
+    tokens = processor.tokenizer.convert_ids_to_tokens(encoding['input_ids'][0])
+
+    text = " ".join([tokens[i] for i in range(len(tokens)) if predicted_class[0][i] != -100])
+    return json.dumps({"filename": os.path.basename(image_path), "content": text})
+
 def convert_file_to_json(file_path: str, file_type: str) -> str:
     try:
         if file_type == "csv":
@@ -42,9 +62,7 @@ def convert_file_to_json(file_path: str, file_type: str) -> str:
             except:
                 df = pd.read_excel(file_path, engine="xlrd", header=None, dtype=str)
         elif file_type == "pdf":
-            with pdfplumber.open(file_path) as pdf:
-                text = "\n".join([page.extract_text() or "" for page in pdf.pages])
-            return json.dumps({"filename": os.path.basename(file_path), "content": text.strip()})
+            return use_layoutlmv3_on_image(file_path)
         else:
             return json.dumps({"error": f"Unsupported file type: {file_type}"})