Merge pull request #381 from alexk307/pdf-ocr

Convert PDF to PNG for OCR
target · Jul 10, 2023 · 0d8cdae · 0d8cdae
2 parents ba48eac + 93474eb
commit 0d8cdae
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/configs/python/backend/backend.yaml b/configs/python/backend/backend.yaml
@@ -445,11 +445,14 @@ scanners:
           - 'image/x-ms-bmp'
           - 'bmp_file'
           - 'image/webp'
+          - 'application/pdf'
+          - 'pdf_file'
       priority: 5
       options:
         extract_text: False
         split_words: True
         tmp_directory: '/dev/shm/'
+        pdf_to_png: True
   'ScanOle':
     - positive:
         flavors:

diff --git a/src/python/strelka/scanners/scan_ocr.py b/src/python/strelka/scanners/scan_ocr.py
@@ -2,6 +2,8 @@
 import subprocess
 import tempfile
 
+import fitz
+
 from strelka import strelka
 
 
@@ -20,6 +22,11 @@ def scan(self, data, file, options, expire_at):
         extract_text = options.get("extract_text", False)
         split_words = options.get("split_words", True)
         tmp_directory = options.get("tmp_directory", "/tmp/")
+        pdf_to_png = options.get("pdf_to_png", False)
+
+        if pdf_to_png and "application/pdf" in file.flavors.get("mime", []):
+            doc = fitz.open(stream=data, filetype="pdf")
+            data = doc.get_page_pixmap(0).tobytes("png")
 
         with tempfile.NamedTemporaryFile(dir=tmp_directory) as tmp_data:
             tmp_data.write(data)