run/ocr/ocr.py

   1 import argparse
   2 import os
   3 from ollama_ocr import OCRProcessor
   4
   5 def ocr_image(ocr, image_path):
   6     """Perform OCR on an image and return the extracted text."""
   7     try:
   8         print(f"Processing {image_path}...")
   9         result = ocr.process_image(
  10                 image_path=image_path,
  11                 format_type='json',
  12                 custom_prompt="Extract all text",
  13                 language="English"
  14         )
  15         return result
  16     except Exception as e:
  17         print(f"Error processing {image_path}: {e}")
  18         return None
  19
  20 def ocr_directory(input_dir):
  21     """Recursively process all files in the directory and its subdirectories."""
  22     ocr = OCRProcessor(model_name='llama3.2-vision:11b')
  23     for root, _, files in os.walk(input_dir):
  24         for file_name in files:
  25             # Check if the file is an image
  26             if file_name.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
  27                 input_path = os.path.join(root, file_name)
  28                 output_path = f'{input_path}.ocr.json'
  29
  30                 text = None
  31                 if os.path.exists(output_path):
  32                     print(f'{output_path} exists.')
  33                 else:
  34                     # Perform OCR on the image
  35                     text = ocr_image(ocr, input_path)
  36                 if text:
  37                     # Write the OCR result to a JSON file
  38                     json_output_path = output_path.rsplit('.', 1)[0]
  39
  40                     with open(json_output_path, 'w') as f:
  41                         f.write(f'{{"text": "{text}"}}')
  42                     print(f"Saved OCR result for {input_path} to {json_output_path}")
  43
  44
  45 if __name__ == "__main__":
  46     parser = argparse.ArgumentParser(description="OCR all images in a directory and save the results as JSON files.")
  47     parser.add_argument('directory', type=str, help='Path to the directory containing image files')
  48
  49     args = parser.parse_args()
  50     input_directory = args.directory
  51     ocr_directory(input_directory)
  52