diff --git a/src/documents/tests/test_zone_ocr.py b/src/documents/tests/test_zone_ocr.py index cfd05af45..81264d829 100644 --- a/src/documents/tests/test_zone_ocr.py +++ b/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() diff --git a/src/documents/zone_ocr.py b/src/documents/zone_ocr.py index 4eec9357f..9d659d490 100644 --- a/src/documents/zone_ocr.py +++ b/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: