Unify OCR zone extraction for preview and saved templates

Co-authored-by: OpenAI Codex <codex@openai.com>

src/documents/tests/test_zone_ocr.py

@@ -7,7 +7,9 @@ from unittest.mock import patch
 
 from django.contrib.auth.models import Group
 from django.contrib.auth.models import User
+from django.test import SimpleTestCase
 from django.test import TestCase
+from PIL import Image
 
 from documents.models import CustomField
 from documents.models import CustomFieldInstance
@@ -17,8 +19,12 @@ from documents.models import OcrTemplate
 from documents.models import OcrTemplateZone
 from documents.zone_ocr import _apply_transform
 from documents.zone_ocr import _convert_value
+from documents.zone_ocr import _crop_zone
 from documents.zone_ocr import _detect_mime
+from documents.zone_ocr import _extract_zone
+from documents.zone_ocr import _process_template
 from documents.zone_ocr import _resolve_doc_path
+from documents.zone_ocr import extract_zone_preview
 from documents.zone_ocr import get_templates_for_document
 from documents.zone_ocr import run_zone_extraction
 from paperless_testing.dirs import DirectoriesMixin
@@ -26,6 +32,171 @@ from paperless_testing.factories import DocumentFactory
 from paperless_testing.permissions import grant_object
 
 
+class TestZoneExtraction(SimpleTestCase):
+    def setUp(self):
+        directory = tempfile.TemporaryDirectory()
+        self.addCleanup(directory.cleanup)
+        self.tmp_path = Path(directory.name)
+        self.page = self.tmp_path / "page.png"
+        with Image.new("RGB", (200, 100), "white") as image:
+            image.paste("black", (20, 10, 60, 30))
+            image.save(self.page)
+        self.zone = OcrTemplateZone(
+            name="Account",
+            x=10,
+            y=5,
+            width=20,
+            height=10,
+            transform="uppercase",
+            ocr_language="eng",
+        )
+        settings = self.settings(SCRATCH_DIR=self.tmp_path)
+        settings.enable()
+        self.addCleanup(settings.disable)
+
+    def test_crop_scales_coordinates(self):
+        with _crop_zone(self.page, self.zone, 100, 50) as cropped:
+            self.assertEqual(cropped.size, (40, 20))
+            self.assertEqual(cropped.getextrema(), ((0, 0), (0, 0), (0, 0)))
+
+    def test_crop_defaults_to_image_dimensions(self):
+        with _crop_zone(self.page, self.zone, 0, 0) as cropped:
+            self.assertEqual(cropped.size, (20, 10))
+            self.assertEqual(cropped.getpixel((0, 0)), (255, 255, 255))
+            self.assertEqual(cropped.getpixel((19, 9)), (0, 0, 0))
+
+    def test_crop_clamps_to_image(self):
+        self.zone.x, self.zone.y = 90, 40
+        with _crop_zone(self.page, self.zone, 100, 50) as cropped:
+            self.assertEqual(cropped.size, (20, 20))
+
+    @patch("documents.zone_ocr.subprocess.run")
+    def test_tiny_crop_skips_reader(self, run):
+        self.zone.width = 0
+        result = _extract_zone(self.page, self.zone, 100, 50, self.tmp_path)
+        self.assertIsNone(result.raw_text)
+        self.assertIsNone(result.value)
+        run.assert_not_called()
+
+    @patch("documents.zone_ocr.subprocess.run")
+    def test_unreadable_image_skips_reader(self, run):
+        self.page.write_bytes(b"not an image")
+        result = _extract_zone(self.page, self.zone, 100, 50, self.tmp_path)
+        self.assertIsNone(result.raw_text)
+        self.assertIsNone(result.value)
+        run.assert_not_called()
+
+    @patch("documents.zone_ocr.subprocess.run")
+    def test_text_extraction_retains_raw_text_and_transforms(self, run):
+        def read_crop(args, **kwargs):
+            with Image.open(args[1]) as crop:
+                self.assertEqual(crop.size, (40, 20))
+                self.assertEqual(crop.getpixel((0, 0)), (0, 0, 0))
+            return MagicMock(stdout="  account abc\n")
+
+        run.side_effect = read_crop
+        result = _extract_zone(self.page, self.zone, 100, 50, self.tmp_path)
+        self.assertEqual(result.raw_text, "account abc")
+        self.assertEqual(result.value, "ACCOUNT ABC")
+        self.assertEqual(run.call_args.args[0][-4:], ["-l", "eng", "--psm", "6"])
+        self.assertEqual(run.call_args.kwargs["timeout"], 30)
+
+    @patch("documents.zone_ocr.subprocess.run")
+    def test_no_text_returns_no_value(self, run):
+        run.return_value.stdout = " \n"
+        result = _extract_zone(self.page, self.zone, 100, 50, self.tmp_path)
+        self.assertIsNone(result.raw_text)
+        self.assertIsNone(result.value)
+
+    @patch("documents.zone_ocr.subprocess.run")
+    def test_date_transform_uses_zone_format(self, run):
+        self.zone.transform = "date"
+        self.zone.date_format = "%d/%m/%Y"
+        run.return_value.stdout = "25/12/2025"
+        result = _extract_zone(self.page, self.zone, 100, 50, self.tmp_path)
+        self.assertEqual(result.raw_text, "25/12/2025")
+        self.assertEqual(result.value, "2025-12-25")
+
+    @patch("documents.zone_ocr.subprocess.run")
+    @patch("zxingcpp.read_barcodes")
+    def test_barcode_extraction_bypasses_tesseract(self, read_barcodes, run):
+        self.zone.transform = "qr_code"
+
+        def read_crop(image):
+            self.assertEqual(image.size, (40, 20))
+            self.assertEqual(image.getpixel((0, 0)), (0, 0, 0))
+            return [MagicMock(text="  https://example.com  ")]
+
+        read_barcodes.side_effect = read_crop
+        result = _extract_zone(self.page, self.zone, 100, 50, self.tmp_path)
+        self.assertEqual(result.raw_text, "  https://example.com  ")
+        self.assertEqual(result.value, "https://example.com")
+        run.assert_not_called()
+        with self.assertRaises(ValueError):
+            read_barcodes.call_args.args[0].getpixel((0, 0))
+
+    @patch("zxingcpp.read_barcodes", return_value=[])
+    def test_missing_barcode_returns_no_value(self, read_barcodes):
+        self.zone.transform = "qr_code"
+        result = _extract_zone(self.page, self.zone, 100, 50, self.tmp_path)
+        self.assertIsNone(result.raw_text)
+        self.assertIsNone(result.value)
+
+    @patch("documents.zone_ocr.subprocess.run")
+    def test_preview_uses_image_dimensions_without_writing(self, run):
+        run.return_value.stdout = "account abc"
+        self.assertEqual(
+            extract_zone_preview(self.page, self.zone, 0, 0, 1),
+            {"raw_text": "account abc", "value": "ACCOUNT ABC"},
+        )
+
+    @patch("documents.zone_ocr.subprocess.run")
+    def test_preview_unavailable_page_skips_reader(self, run):
+        self.zone.page = 2
+        self.assertEqual(
+            extract_zone_preview(self.page, self.zone, 100, 50, 1),
+            {"raw_text": None, "value": None},
+        )
+        run.assert_not_called()
+
+
+class TestSavedZoneExtraction(DirectoriesMixin, TestCase):
+    @patch("documents.zone_ocr.subprocess.run")
+    def test_saved_extraction_matches_preview_and_rejects_invalid_value(self, run):
+        page = Path(self.dirs.scratch_dir) / "page.png"
+        with Image.new("RGB", (200, 100), "white") as image:
+            image.save(page)
+        document = DocumentFactory(title="Original")
+        template = OcrTemplate.objects.create(
+            name="Account",
+            source_width=100,
+            source_height=50,
+        )
+        zone = OcrTemplateZone.objects.create(
+            template=template,
+            name="Title",
+            target="title",
+            x=0,
+            y=0,
+            width=20,
+            height=10,
+            transform="uppercase",
+            validation_regex="ACCOUNT [A-Z]+",
+        )
+        for raw_text, expected_title in [
+            ("account abc", "ACCOUNT ABC"),
+            ("invalid", "ACCOUNT ABC"),
+        ]:
+            with self.subTest(raw_text=raw_text):
+                run.return_value.stdout = raw_text
+                preview = extract_zone_preview(page, zone, 100, 50, 1)
+                self.assertEqual(preview["raw_text"], raw_text)
+                self.assertEqual(preview["value"], raw_text.upper())
+                _process_template(document, page, template, [zone])
+                document.refresh_from_db()
+                self.assertEqual(document.title, expected_title)
+
+
 class TestTemplateVisibility(DirectoriesMixin, TestCase):
     def setUp(self):
         super().setUp()

src/documents/zone_ocr.py

@@ -14,6 +14,7 @@ import re
 import string
 import subprocess
 import tempfile
+from dataclasses import dataclass
 from datetime import date
 from datetime import datetime
 from pathlib import Path
@@ -34,6 +35,12 @@ from documents.versioning import get_root_document
 logger = logging.getLogger("paperless.zone_ocr")
 
 
+@dataclass(frozen=True)
+class ZoneExtraction:
+    raw_text: str | None = None
+    value: str | None = None
+
+
 def get_templates_for_document(document: Document) -> list[OcrTemplate]:
     """
     The enabled templates that match the document's content and that its owner
@@ -171,13 +178,14 @@ def _process_template(
             src_w = zone.zone_source_width or template.source_width
             src_h = zone.zone_source_height or template.source_height
 
-            extracted = _extract_zone(
+            result = _extract_zone(
                 page_images[page_idx],
                 zone,
                 src_w,
                 src_h,
                 tmp_path,
             )
+            extracted = result.value
 
             if (
                 extracted is not None
@@ -323,12 +331,13 @@ def _crop_zone(
     zone: OcrTemplateZone,
     source_width: int,
     source_height: int,
-    tmp_dir: Path,
 ) -> Image.Image | None:
     """Crop a zone from the page image and return the PIL Image."""
     try:
         with Image.open(page_img) as img:
             img_width, img_height = img.size
+            if not source_width or not source_height:
+                source_width, source_height = img.size
 
             scale_x = img_width / source_width
             scale_y = img_height / source_height
@@ -348,7 +357,7 @@ def _crop_zone(
                 logger.warning("Zone OCR: crop too small for zone '%s'", zone.name)
                 return None
 
-            return img.crop((crop_left, crop_top, crop_right, crop_bottom)).copy()
+            return img.crop((crop_left, crop_top, crop_right, crop_bottom))
     except Exception:
         logger.exception("Zone OCR: crop failed for zone '%s'", zone.name)
         return None
@@ -421,31 +430,25 @@ def _extract_zone(
     source_width: int,
     source_height: int,
     tmp_dir: Path,
-) -> str | None:
-    """Crop a zone from the page image and extract text via OCR or barcode reader."""
-    cropped = _crop_zone(page_img, zone, source_width, source_height, tmp_dir)
+) -> ZoneExtraction:
+    """Read and transform a zone, retaining the text for the editor's preview."""
+    cropped = _crop_zone(page_img, zone, source_width, source_height)
     if cropped is None:
-        return None
+        return ZoneExtraction()
 
-    # QR/barcode zones skip Tesseract entirely
-    if zone.transform == "qr_code":
-        text = _read_barcode(cropped, zone.name)
-        if not text:
-            return None
-        return _apply_transform(
-            text,
-            zone.transform,
-            getattr(zone, "date_format", "") or "",
+    try:
+        raw_text = (
+            _read_barcode(cropped, zone.name)
+            if zone.transform == OcrTemplateZone.TransformType.QR_CODE
+            else _ocr_text(cropped, zone, tmp_dir)
         )
-
-    text = _ocr_text(cropped, zone, tmp_dir)
-    if not text:
-        return None
-
-    return _apply_transform(
-        text,
-        zone.transform,
-        getattr(zone, "date_format", "") or "",
+    finally:
+        cropped.close()
+    return ZoneExtraction(
+        raw_text=raw_text,
+        value=_apply_transform(raw_text, zone.transform, zone.date_format)
+        if raw_text
+        else None,
     )
 
 
@@ -456,15 +459,7 @@ def extract_zone_preview(
     source_height: int,
     page_count: int | None,
 ) -> dict:
-    """Non-destructive single-zone extraction for the editor's per-zone test.
-
-    Renders the zone's page, crops it, runs OCR (or the barcode reader) and
-    applies the transform — WITHOUT writing any custom field. Returns the raw
-    OCR text and the transformed value so the user can see what the zone yields
-    (and tune the validation regex) before saving.
-    """
-    # zone.page is 1-indexed (1 = first, -1 = last); resolve to a 0-indexed
-    # image index exactly like the production extraction path does.
+    """Extract a zone without writing its value to the document."""
     page_idx = _resolve_page_idx(zone.page, page_count)
     with tempfile.TemporaryDirectory(dir=settings.SCRATCH_DIR) as tmp_dir:
         tmp_path = Path(tmp_dir)
@@ -472,35 +467,14 @@ def extract_zone_preview(
         if page_idx not in page_images:
             return {"raw_text": None, "value": None}
 
-        if not source_width or not source_height:
-            with Image.open(page_images[page_idx]) as im:
-                source_width, source_height = im.size
-
-        cropped = _crop_zone(
+        result = _extract_zone(
             page_images[page_idx],
             zone,
             source_width,
             source_height,
             tmp_path,
         )
-        if cropped is None:
-            return {"raw_text": None, "value": None}
-
-        if zone.transform == "qr_code":
-            raw_text = _read_barcode(cropped, zone.name)
-        else:
-            raw_text = _ocr_text(cropped, zone, tmp_path)
-
-        value = (
-            _apply_transform(
-                raw_text,
-                zone.transform,
-                getattr(zone, "date_format", "") or "",
-            )
-            if raw_text
-            else None
-        )
-        return {"raw_text": raw_text, "value": value}
+        return {"raw_text": result.raw_text, "value": result.value}
 
 
 def _parse_date(text: str, fmt: str) -> str: