Simple OCR script for Ollama

diff --git a/files/ocr/ocr.py b/files/ocr/ocr.py
new file mode 100644 (file)
index 0000000..795a83b
--- /dev/null
+++ b/files/ocr/ocr.py
@@ -0,0 +1,49 @@
+import argparse
+import os
+from ollama_ocr import OCRProcessor
+
+def ocr_image(ocr, image_path):
+    """Perform OCR on an image and return the extracted text."""
+    try:
+        print(f"Processing {image_path}...")
+        result = ocr.process_image(
+                image_path=image_path,
+                format_type='json',
+                custom_prompt="Extract all text",
+                language="English"
+        )
+        return result
+    except Exception as e:
+        print(f"Error processing {image_path}: {e}")
+        return None
+
+def ocr_directory(input_dir):
+    """Recursively process all files in the directory and its subdirectories."""
+    ocr = OCRProcessor(model_name='llama3.2-vision:11b')
+    for root, _, files in os.walk(input_dir):
+        for file_name in files:
+            # Check if the file is an image
+            if file_name.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
+                input_path = os.path.join(root, file_name)
+                output_path = os.path.join(root, file_name.replace('.', '_ocr.'))
+
+                # Perform OCR on the image
+                text = ocr_image(ocr, input_path)
+
+                if text:
+                    # Write the OCR result to a JSON file
+                    json_output_path = output_path.rsplit('.', 1)[0] + '.json'
+
+                    with open(json_output_path, 'w') as f:
+                        f.write(f'{{"text": "{text}"}}')
+                    print(f"Saved OCR result for {input_path} to {json_output_path}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="OCR all images in a directory and save the results as JSON files.")
+    parser.add_argument('directory', type=str, help='Path to the directory containing image files')
+
+    args = parser.parse_args()
+    input_directory = args.directory
+    ocr_directory(input_directory)
+

diff --git a/files/ocr/requirements.txt b/files/ocr/requirements.txt
new file mode 100644 (file)
index 0000000..02f390c
--- /dev/null
+++ b/files/ocr/requirements.txt
@@ -0,0 +1,2 @@
+ollama-ocr
+# https://pypi.org/project/ollama-ocr/

diff --git a/files/ocr/run.sh b/files/ocr/run.sh
new file mode 100755 (executable)
index 0000000..0a02fa7
--- /dev/null
+++ b/files/ocr/run.sh
@@ -0,0 +1,2 @@
+echo "Processing... $IMG_DIR"
+$HOME/envs/ocr/bin/python ocr.py $IMG_DIR