Merge pull request NanoNets#23 from sirius116/enhancement/pdf_support

PDF Support, Multi File Support and Some Fixes

Author: mandalsouvik3333

docext/app/app.py

@@ -14,6 +14,7 @@
 from docext.core.config import TEMPLATES_FIELDS
 from docext.core.config import TEMPLATES_TABLES
 from docext.core.extract import extract_information
+from docext.core.utils import convert_files_to_images
 from docext.core.vllm import VLLMServer
 
 METADATA = []
@@ -139,8 +140,28 @@ def define_keys_and_extract(model_name: str, max_img_size: int, concurrency_limi
                 visible=False,
             )
 
-            images_input = gr.Gallery(label="Upload images", preview=True)
-            submit_btn = gr.Button("Submit")
+            file_input = gr.File(
+                label="Upload Documents",
+                file_types=[".pdf", ".jpg", ".jpeg", ".png", ".tiff", ".bmp", ".gif", ".webp"],
+                file_count="multiple"
+            )
+            images_input = gr.Gallery(label="Document Preview", preview=True, visible=False)
+            submit_btn = gr.Button("Submit", visible=False)
+
+            def handle_file_upload(files):
+                if not files:
+                    return None, gr.update(visible=False), gr.update(visible=False)
+                
+                file_paths = [f.name for f in files]
+                # Convert PDFs to images if necessary and get all image paths
+                image_paths = convert_files_to_images(file_paths)
+                return image_paths, gr.update(visible=True, value=image_paths), gr.update(visible=True)
+
+            file_input.change(
+                handle_file_upload,
+                inputs=[file_input],
+                outputs=[images_input, images_input, submit_btn]
+            )
 
     with gr.Row():
         with gr.Column(scale=3):

docext/core/client.py

@@ -42,7 +42,9 @@ def sync_request(
     elif model_name.startswith("openrouter"):
         completion_args["response_format"] = format
     elif "gpt" in model_name.lower():
-        completion_args["response_format"] = {"type": "json_object"}
+        # Only set response_format if the prompt mentions "json"
+        if any("json" in m.get("text", "").lower() for m in messages if isinstance(m, dict)):
+            completion_args["response_format"] = {"type": "json_object"}
 
     response = completion(**completion_args)
     return response.json()

docext/core/extract.py

@@ -13,6 +13,7 @@
 from docext.core.confidence import get_fields_confidence_score_messages_binary
 from docext.core.prompts import get_fields_messages
 from docext.core.prompts import get_tables_messages
+from docext.core.utils import convert_files_to_images
 from docext.core.utils import resize_images
 from docext.core.utils import validate_fields_and_tables
 from docext.core.utils import validate_file_paths
@@ -65,14 +66,36 @@ def extract_fields_from_documents(
     extracted_fields = json_repair.loads(response)
     conf_scores = json_repair.loads(response_conf_score)
 
-    df = pd.DataFrame(
-        {
-            "fields": field_names,
-            "answer": [extracted_fields.get(field, "") for field in field_names],
-            "confidence": [conf_scores.get(field, "Low") for field in field_names],
-        },
-    )
-    return df
+    logger.info(f"Extracted fields: {extracted_fields}")
+    logger.info(f"Conf scores: {conf_scores}")
+
+    # Handle both single dictionary and list of dictionaries
+    if not isinstance(extracted_fields, list):
+        extracted_fields = [extracted_fields]
+    
+    # Handle confidence scores similarly
+    if not isinstance(conf_scores, list):
+        conf_scores = [conf_scores] * len(extracted_fields)
+    elif len(conf_scores) < len(extracted_fields):
+        # If we have fewer confidence scores than documents, pad with the first confidence score
+        conf_scores.extend([conf_scores[0]] * (len(extracted_fields) - len(conf_scores)))
+    
+    # Create a list of dataframes, one for each document
+    dfs = []
+    for idx, (doc_fields, doc_conf_scores) in enumerate(zip(extracted_fields, conf_scores)):
+        df = pd.DataFrame(
+            {
+                "fields": field_names,
+                "answer": [doc_fields.get(field, "") for field in field_names],
+                "confidence": [doc_conf_scores.get(field, "Low") for field in field_names],
+                "document_index": [idx] * len(field_names)
+            },
+        )
+        dfs.append(df)
+    
+    # Concatenate all dataframes with a document index
+    final_df = pd.concat(dfs, ignore_index=True)
+    return final_df
 
 
 def extract_tables_from_documents(
@@ -112,6 +135,7 @@ def extract_information(
         for file_input in file_inputs
     ]
     validate_file_paths(file_paths)
+    file_paths = convert_files_to_images(file_paths)
     resize_images(file_paths, max_img_size)
 
     # call fields and tables extraction in parallel
@@ -131,4 +155,9 @@ def extract_information(
 
         fields_df = future_fields.result()
         tables_df = future_tables.result()
+    
+    # Group fields by document_index for better display
+    if not fields_df.empty and 'document_index' in fields_df.columns:
+        fields_df = fields_df.sort_values(['document_index', 'fields'])
+    
     return fields_df, tables_df

docext/core/file_converters/file_converter.py

@@ -0,0 +1,10 @@
+from __future__ import annotations
+
+from abc import ABC
+from abc import abstractmethod
+
+
+class FileConverter(ABC):
+    @abstractmethod
+    def convert_to_images(self, file_path: str):
+        pass

docext/core/file_converters/pdf_converter.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+from docext.core.file_converters.file_converter import FileConverter
+from pdf2image import convert_from_path
+
+
+class PDFConverter(FileConverter):
+    def convert_to_images(self, file_path: str):
+        return convert_from_path(file_path)

docext/core/utils.py

@@ -7,6 +7,7 @@
 
 import pandas as pd
 from PIL import Image
+from docext.core.file_converters.pdf_converter import PDFConverter
 
 
 def encode_image(image_path):
@@ -65,4 +66,31 @@ def validate_file_paths(file_paths: list[str]):
             ".bmp",
             ".gif",
             ".webp",
+            ".pdf",
         ], f"File {file_path} is not an image"
+
+def file_is_supported_image(file_path: str) -> bool:
+    return os.path.splitext(file_path)[1].lower() in [
+        ".jpg",
+        ".jpeg",
+        ".png",
+        ".tiff",
+        ".bmp",
+        ".gif",
+        ".webp",
+    ]
+
+# TODO: add support for other file types; only support pdf for now
+def convert_files_to_images(file_paths: list[str]):
+    converted_file_paths = []
+    pdf_converter = PDFConverter()
+    for file_path in file_paths:
+        if os.path.splitext(file_path)[1].lower() == ".pdf":
+            images = pdf_converter.convert_to_images(file_path)
+            for i, image in enumerate(images):
+                image.save(f"{file_path.replace('.pdf', '')}_{i}.jpg")
+                converted_file_paths.append(f"{file_path.replace('.pdf', '')}_{i}.jpg")
+        else:
+            if file_is_supported_image(file_path):
+                converted_file_paths.append(file_path)
+    return converted_file_paths

requirements.txt

@@ -16,3 +16,4 @@ tenacity
 types-requests
 vllm==v0.8.3
 xgrammar==0.1.17
+pdf2image