vlm4ocr 0.2.0__tar.gz → 0.3.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vlm4ocr
-Version: 0.2.0
+Version: 0.3.1
 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.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "vlm4ocr"
-version = "0.2.0"
+version = "0.3.1"
 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.1/vlm4ocr/__init__.py

@@ -0,0 +1,15 @@
+from .ocr_engines import OCREngine
+from .vlm_engines import BasicVLMConfig, ReasoningVLMConfig, OpenAIReasoningVLMConfig, OllamaVLMEngine, OpenAICompatibleVLMEngine, VLLMVLMEngine, OpenRouterVLMEngine, OpenAIVLMEngine, AzureOpenAIVLMEngine
+
+__all__ = [
+    "BasicVLMConfig",
+    "ReasoningVLMConfig", 
+    "OpenAIReasoningVLMConfig",
+    "OCREngine",
+    "OllamaVLMEngine",
+    "OpenAICompatibleVLMEngine",
+    "VLLMVLMEngine",
+    "OpenRouterVLMEngine",
+    "OpenAIVLMEngine",
+    "AzureOpenAIVLMEngine"
+]

vlm4ocr-0.3.1/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.1}/vlm4ocr/cli.py

@@ -4,18 +4,9 @@ import sys
 import logging
 import asyncio
 import time
-
-# Attempt to import from the local package structure
-try:
-    from .ocr_engines import OCREngine
-    from .vlm_engines import OpenAIVLMEngine, AzureOpenAIVLMEngine, OllamaVLMEngine, BasicVLMConfig
-    from .data_types import OCRResult
-except ImportError:
-    # Fallback for when the package is installed
-    from vlm4ocr.ocr_engines import OCREngine
-    from vlm4ocr.vlm_engines import OpenAIVLMEngine, AzureOpenAIVLMEngine, OllamaVLMEngine, BasicVLMConfig
-    from vlm4ocr.data_types import OCRResult
-
+from .ocr_engines import OCREngine
+from .vlm_engines import OpenAICompatibleVLMEngine, OpenAIVLMEngine, AzureOpenAIVLMEngine, OllamaVLMEngine, BasicVLMConfig
+from .data_types import OCRResult
 import tqdm.asyncio
 
 # --- Global logger setup (console) ---
@@ -208,7 +199,7 @@ def main():
             vlm_engine_instance = OpenAIVLMEngine(model=args.model, api_key=args.api_key, config=config)
         elif args.vlm_engine == "openai_compatible":
             if not args.base_url: parser.error("--base_url is required for openai_compatible.")
-            vlm_engine_instance = OpenAIVLMEngine(model=args.model, api_key=args.api_key, base_url=args.base_url, config=config)
+            vlm_engine_instance = OpenAICompatibleVLMEngine(model=args.model, api_key=args.api_key, base_url=args.base_url, config=config)
         elif args.vlm_engine == "azure_openai":
             if not args.azure_api_key: parser.error("--azure_api_key (or AZURE_OPENAI_API_KEY) is required.")
             if not args.azure_endpoint: parser.error("--azure_endpoint (or AZURE_OPENAI_ENDPOINT) is required.")

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

@@ -1,9 +1,9 @@
 import os
-from typing import List, Literal
+from typing import List, Dict, 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:
@@ -24,6 +24,7 @@ class OCRResult:
     pages: List[dict] = field(default_factory=list)
     filename: str = field(init=False)
     status: str = field(init=False, default="processing")
+    messages_log: List[List[Dict[str,str]]] = field(default_factory=list)
 
     def __post_init__(self):
         """
@@ -33,8 +34,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):
@@ -67,10 +68,6 @@ class OCRResult:
         }
         self.pages.append(page)
 
-
-    def __len__(self):
-        return len(self.pages)
-
     def get_page(self, idx):
         if not isinstance(idx, int):
             raise ValueError("Index must be an integer")
@@ -78,6 +75,21 @@ class OCRResult:
             raise IndexError(f"Index out of range. The OCRResult has {len(self.pages)} pages, but index {idx} was requested.")
         
         return self.pages[idx]
+
+    def clear_messages_log(self):
+        self.messages_log = []
+
+    def add_messages_to_log(self, messages: List[Dict[str,str]]):
+        if not isinstance(messages, list):
+            raise ValueError("messages must be a list of dict")
+        
+        self.messages_log.extend(messages)
+
+    def get_messages_log(self) -> List[List[Dict[str,str]]]:
+        return self.messages_log.copy()
+
+    def __len__(self):
+        return len(self.pages)
     
     def __iter__(self):
         return iter(self.pages)

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

@@ -3,10 +3,10 @@ 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
+from vlm4ocr.vlm_engines import VLMEngine, MessagesLogger
 
 SUPPORTED_IMAGE_EXTS = ['.pdf', '.tif', '.tiff', '.png', '.jpg', '.jpeg', '.bmp', '.gif', '.webp']
 
@@ -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()
@@ -123,7 +126,8 @@ class OCREngine:
                     stream=True
                 )
                 for chunk in response_stream:
-                    yield {"type": "ocr_chunk", "data": chunk}
+                    if chunk["type"] == "response":
+                        yield {"type": "ocr_chunk", "data": chunk["data"]}
 
                 if i < len(images) - 1:
                     yield {"type": "page_delimiter", "data": get_default_page_delimiter(self.output_mode)}
@@ -154,7 +158,8 @@ class OCREngine:
                     stream=True
                 )
             for chunk in response_stream:
-                yield {"type": "ocr_chunk", "data": chunk}
+                if chunk["type"] == "response":
+                    yield {"type": "ocr_chunk", "data": chunk["data"]}
             
 
     def sequential_ocr(self, file_paths: Union[str, Iterable[str]], rotate_correction:bool=False, 
@@ -268,18 +273,32 @@ class OCREngine:
 
                 try:
                     messages = self.vlm_engine.get_ocr_messages(self.system_prompt, self.user_prompt, image)
+                    # Define a messages logger to capture messages
+                    messages_logger = MessagesLogger()
+                    # Generate response
                     response = self.vlm_engine.chat(
                         messages,
                         verbose=verbose,
-                        stream=False
+                        stream=False,
+                        messages_logger=messages_logger
                     )
+                    ocr_text = response["response"]
                     # Clean the response if output mode is markdown
                     if self.output_mode == "markdown":
-                        response = clean_markdown(response)
+                        ocr_text = clean_markdown(ocr_text)
+
+                    # Parse the response if output mode is JSON
+                    elif 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)
                     
                     # Add the page to the OCR result
-                    ocr_result.add_page(text=response, 
+                    ocr_result.add_page(text=ocr_text, 
                                         image_processing_status=image_processing_status)
+                    
+                    # Add messages log to the OCR result
+                    ocr_result.add_messages_to_log(messages_logger.get_messages_log())
                 
                 except Exception as page_e:
                     ocr_result.status = "error"
@@ -378,6 +397,7 @@ class OCREngine:
             filename = os.path.basename(file_path)
             file_ext = os.path.splitext(file_path)[1].lower()
             result = OCRResult(input_dir=file_path, output_mode=self.output_mode)
+            messages_logger = MessagesLogger()
             # check file extension
             if file_ext not in SUPPORTED_IMAGE_EXTS:
                 result.status = "error"
@@ -407,7 +427,8 @@ class OCREngine:
                         data_loader=data_loader,
                         page_index=page_index,
                         rotate_correction=rotate_correction,
-                        max_dimension_pixels=max_dimension_pixels
+                        max_dimension_pixels=max_dimension_pixels,
+                        messages_logger=messages_logger
                     )
                     page_processing_tasks.append(task)
                 
@@ -419,14 +440,17 @@ class OCREngine:
             except Exception as e:
                 result.status = "error"
                 result.add_page(text=f"Error during OCR for {filename}: {str(e)}", image_processing_status={})
+                result.add_messages_to_log(messages_logger.get_messages_log())
                 return result
 
         # Set status to success if no errors occurred
         result.status = "success"
+        result.add_messages_to_log(messages_logger.get_messages_log())
         return result
 
     async def _ocr_page_with_semaphore(self, vlm_call_semaphore: asyncio.Semaphore, data_loader: DataLoader,
-                                       page_index:int, rotate_correction:bool=False, max_dimension_pixels:int=None) -> Tuple[str, Dict[str, str]]:
+                                       page_index:int, rotate_correction:bool=False, max_dimension_pixels:int=None,
+                                       messages_logger:MessagesLogger=None) -> Tuple[str, Dict[str, str]]:
         """
         This internal method takes a semaphore and OCR a single image/page using the VLM inference engine.
 
@@ -467,9 +491,19 @@ class OCREngine:
                     }
 
             messages = self.vlm_engine.get_ocr_messages(self.system_prompt, self.user_prompt, image)
-            ocr_text = await self.vlm_engine.chat_async( 
+            response = await self.vlm_engine.chat_async( 
                 messages,
+                messages_logger=messages_logger
             )
+            ocr_text = response["response"]
+            # 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
+            elif 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.1}/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: