vlm4ocr 0.2.0__tar.gz → 0.3.0__tar.gz

{vlm4ocr-0.2.0 → vlm4ocr-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vlm4ocr
-Version: 0.2.0
+Version: 0.3.0
 Summary: Python package and Web App for OCR with vision language models.
 License: MIT
 Author: Enshuo (David) Hsu
@@ -10,6 +10,8 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Provides-Extra: tesseract
+Requires-Dist: colorama (>=0.4.4)
+Requires-Dist: json-repair (>=0.30.0)
 Requires-Dist: pdf2image (>=1.16.0)
 Requires-Dist: pillow (>=10.0.0)
 Requires-Dist: pytesseract (>=0.3.13) ; extra == "tesseract"

{vlm4ocr-0.2.0 → vlm4ocr-0.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "vlm4ocr"
-version = "0.2.0"
+version = "0.3.0"
 description = "Python package and Web App for OCR with vision language models."
 authors = ["Enshuo (David) Hsu"]
 license = "MIT"
@@ -15,7 +15,9 @@ exclude = [
 [tool.poetry.dependencies]
 python = "^3.11"
 pdf2image = ">=1.16.0"
+colorama = ">=0.4.4"
 pillow = ">=10.0.0"
+json-repair = ">=0.30.0"
 pytesseract = { version = ">=0.3.13", optional = true }
 
 [tool.poetry.scripts]

vlm4ocr-0.3.0/vlm4ocr/assets/default_prompt_templates/ocr_JSON_system_prompt.txt

@@ -0,0 +1 @@
+You are a helpful assistant that can convert scanned documents into JSON format. Your output is accurate and well-formatted, starting with ```json and ending with ```. You will only output the JSON text without any additional explanations or comments. The JSON should include all text, tables, and lists with appropriate keys and values. You will ignore images, icons, or anything that can not be converted into text.

{vlm4ocr-0.2.0 → vlm4ocr-0.3.0}/vlm4ocr/data_types.py

@@ -3,7 +3,7 @@ from typing import List, Literal
 from dataclasses import dataclass, field
 from vlm4ocr.utils import get_default_page_delimiter
 
-OutputMode = Literal["markdown", "HTML", "text"]
+OutputMode = Literal["markdown", "HTML", "text", "JSON"]
 
 @dataclass
 class OCRResult:
@@ -33,8 +33,8 @@ class OCRResult:
         self.filename = os.path.basename(self.input_dir)
 
         # output_mode validation
-        if self.output_mode not in ["markdown", "HTML", "text"]:
-            raise ValueError("output_mode must be 'markdown', 'HTML', or 'text'")
+        if self.output_mode not in ["markdown", "HTML", "text", "JSON"]:
+            raise ValueError("output_mode must be 'markdown', 'HTML', 'text', or 'JSON'")
 
         # pages validation 
         if not isinstance(self.pages, list):

{vlm4ocr-0.2.0 → vlm4ocr-0.3.0}/vlm4ocr/ocr_engines.py

@@ -3,8 +3,8 @@ from typing import Tuple, List, Dict, Union, Generator, AsyncGenerator, Iterable
 import importlib
 import asyncio
 from colorama import Fore, Style   
-from PIL import Image
-from vlm4ocr.utils import DataLoader, PDFDataLoader, TIFFDataLoader, ImageDataLoader, ImageProcessor, clean_markdown, get_default_page_delimiter
+import json
+from vlm4ocr.utils import DataLoader, PDFDataLoader, TIFFDataLoader, ImageDataLoader, ImageProcessor, clean_markdown, extract_json, get_default_page_delimiter
 from vlm4ocr.data_types import OCRResult
 from vlm4ocr.vlm_engines import VLMEngine
 
@@ -21,7 +21,7 @@ class OCREngine:
         inference_engine : InferenceEngine
             The inference engine to use for OCR.
         output_mode : str, Optional
-            The output format. Must be 'markdown', 'HTML', or 'text'.
+            The output format. Must be 'markdown', 'HTML', 'text', or 'JSON'.
         system_prompt : str, Optional
             Custom system prompt. We recommend use a default system prompt by leaving this blank. 
         user_prompt : str, Optional
@@ -33,8 +33,8 @@ class OCREngine:
         self.vlm_engine = vlm_engine
 
         # Check output mode
-        if output_mode not in ["markdown", "HTML", "text"]:
-            raise ValueError("output_mode must be 'markdown', 'HTML', or 'text'")
+        if output_mode not in ["markdown", "HTML", "text", "JSON"]:
+            raise ValueError("output_mode must be 'markdown', 'HTML', 'text', or 'JSON'.")
         self.output_mode = output_mode
 
         # System prompt
@@ -49,6 +49,9 @@ class OCREngine:
         if isinstance(user_prompt, str) and user_prompt:
             self.user_prompt = user_prompt
         else:
+            if self.output_mode == "JSON":
+                raise ValueError("user_prompt must be provided when output_mode is 'JSON' to define the JSON structure.")
+
             prompt_template_path = importlib.resources.files('vlm4ocr.assets.default_prompt_templates').joinpath(f'ocr_{self.output_mode}_user_prompt.txt')
             with prompt_template_path.open('r', encoding='utf-8') as f:
                 self.user_prompt =  f.read()
@@ -276,6 +279,12 @@ class OCREngine:
                     # Clean the response if output mode is markdown
                     if self.output_mode == "markdown":
                         response = clean_markdown(response)
+
+                    # Parse the response if output mode is JSON
+                    if self.output_mode == "JSON":
+                        json_list = extract_json(response)
+                        # Serialize the JSON list to a string
+                        response = json.dumps(json_list, indent=4)
                     
                     # Add the page to the OCR result
                     ocr_result.add_page(text=response, 
@@ -470,6 +479,14 @@ class OCREngine:
             ocr_text = await self.vlm_engine.chat_async( 
                 messages,
             )
+            # Clean the OCR text if output mode is markdown
             if self.output_mode == "markdown":
                 ocr_text = clean_markdown(ocr_text)
+
+            # Parse the response if output mode is JSON
+            if self.output_mode == "JSON":
+                json_list = extract_json(ocr_text)
+                # Serialize the JSON list to a string
+                ocr_text = json.dumps(json_list, indent=4)
+
             return ocr_text, image_processing_status

{vlm4ocr-0.2.0 → vlm4ocr-0.3.0}/vlm4ocr/utils.py

@@ -2,11 +2,14 @@ import abc
 import os
 import io
 import base64
-from typing import Union, List, Tuple
+from typing import Dict, List, Tuple
+import json
+import json_repair
 import importlib.util
 from pdf2image import convert_from_path, pdfinfo_from_path
 from PIL import Image
 import asyncio
+import warnings
 
 
 class DataLoader(abc.ABC):
@@ -229,6 +232,55 @@ def clean_markdown(text:str) -> str:
     cleaned_text = text.replace("```markdown", "").replace("```", "")
     return cleaned_text
 
+def _find_dict_strings( text: str) -> List[str]:
+    """
+    Extracts balanced JSON-like dictionaries from a string, even if nested.
+
+    Parameters:
+    -----------
+    text : str
+        the input text containing JSON-like structures.
+
+    Returns : List[str]
+        A list of valid JSON-like strings representing dictionaries.
+    """
+    open_brace = 0
+    start = -1
+    json_objects = []
+
+    for i, char in enumerate(text):
+        if char == '{':
+            if open_brace == 0:
+                # start of a new JSON object
+                start = i 
+            open_brace += 1
+        elif char == '}':
+            open_brace -= 1
+            if open_brace == 0 and start != -1:
+                json_objects.append(text[start:i + 1])
+                start = -1
+
+    return json_objects
+
+def extract_json(gen_text:str) -> List[Dict[str, str]]:
+    """ 
+    This method inputs a generated text and output a JSON of information tuples
+    """
+    out = []
+    dict_str_list = _find_dict_strings(gen_text)
+    for dict_str in dict_str_list:
+        try:
+            dict_obj = json.loads(dict_str)
+            out.append(dict_obj)
+        except json.JSONDecodeError:
+            dict_obj = json_repair.repair_json(dict_str, skip_json_loads=True, return_objects=True)
+            if dict_obj:
+                warnings.warn(f'JSONDecodeError detected, fixed with repair_json:\n{dict_str}', RuntimeWarning)
+                out.append(dict_obj)
+            else:
+                warnings.warn(f'JSONDecodeError could not be fixed:\n{dict_str}', RuntimeWarning)
+    return out
+
 def get_default_page_delimiter(output_mode:str) -> str:
     """ 
     Returns the default page delimiter based on the environment variable.
@@ -243,8 +295,8 @@ def get_default_page_delimiter(output_mode:str) -> str:
     str
         The default page delimiter.
     """
-    if output_mode not in ["markdown", "HTML", "text"]:
-        raise ValueError("output_mode must be 'markdown', 'HTML', or 'text'")
+    if output_mode not in ["markdown", "HTML", "text", "JSON"]:
+        raise ValueError("output_mode must be 'markdown', 'HTML', 'text', or 'JSON'")
     
     if output_mode == "markdown":
         return "\n\n---\n\n"
@@ -252,6 +304,8 @@ def get_default_page_delimiter(output_mode:str) -> str:
         return "<br><br>"
     elif output_mode == "text":
         return "\n\n---\n\n"
+    elif output_mode == "JSON":
+        return "\n\n---\n\n"
 
 
 class ImageProcessor: