diff --git a/src/odoo_ocr/pipeline/classifier.py b/src/odoo_ocr/pipeline/classifier.py index 74b2af5..14e9885 100644 --- a/src/odoo_ocr/pipeline/classifier.py +++ b/src/odoo_ocr/pipeline/classifier.py @@ -9,6 +9,7 @@ import fitz # pymupdf from PIL import Image, ImageStat from odoo_ocr.clients import BaseVLMClient +from odoo_ocr.clients.base import load_prompt from odoo_ocr.config import Settings from odoo_ocr.schemas.document import ClassificationResult, DocumentClass @@ -89,7 +90,12 @@ async def classify_document( logger.info("Heuristic classifier: %s", heuristic.category) return heuristic - system_prompt = _load_prompt("classifier_system.txt") + if not pages: + raise ValueError( + f"Document {path} has no renderable pages; cannot classify with the VLM fallback." + ) + + system_prompt = load_prompt("classifier_system.txt") user_prompt = ( "Classify this invoice document. Return only the requested JSON object." ) @@ -104,9 +110,3 @@ async def classify_document( result = client.parse_json(response.content, ClassificationResult) logger.info("VLM classifier: %s (confidence %.2f)", result.category, result.confidence) return result - - -def _load_prompt(name: str) -> str: - from odoo_ocr.clients.base import BaseVLMClient - - return BaseVLMClient._load_prompt_file(Path("prompts") / name) diff --git a/src/odoo_ocr/pipeline/handwritten.py b/src/odoo_ocr/pipeline/handwritten.py index fd3fb0c..9e9b191 100644 --- a/src/odoo_ocr/pipeline/handwritten.py +++ b/src/odoo_ocr/pipeline/handwritten.py @@ -10,7 +10,9 @@ import numpy as np from PIL import Image from odoo_ocr.clients import BaseVLMClient +from odoo_ocr.clients.base import load_prompt from odoo_ocr.config import Settings +from odoo_ocr.schemas import OcrResponse logger = logging.getLogger(__name__) @@ -28,14 +30,6 @@ def _preprocess_handwritten(image: Image.Image) -> Image.Image: return Image.fromarray(processed, mode="L").convert("RGB") -def _load_prompt(name: str) -> str: - from pathlib import Path - - from odoo_ocr.clients.base import BaseVLMClient - - return BaseVLMClient._load_prompt_file(Path("prompts") / name) - - async def process_handwritten( path: Path, pages: list[Image.Image], @@ -43,7 +37,7 @@ async def process_handwritten( settings: Settings, ) -> str: """Run OCR on handwritten invoice pages using the configured handwriting model.""" - system_prompt = _load_prompt("ocr_system.txt") + system_prompt = load_prompt("ocr_system.txt") user_prompt = ( "Read all text from this handwritten invoice accurately. " "Preserve line breaks and mark unclear words with [UNCLEAR]." @@ -56,17 +50,12 @@ async def process_handwritten( system_prompt=system_prompt, user_prompt=f"{user_prompt} (page {i + 1} of {len(pages)})", images=[preprocessed], + response_format=OcrResponse, temperature=0.1, model=settings.models.ocr, ) - import json - - try: - parsed = json.loads(response.content.strip().strip("`").removeprefix("json").strip()) - text = parsed.get("text", response.content) - except Exception: - text = response.content - page_texts.append(f"--- Page {i + 1} ---\n{text.strip()}") + ocr = client.parse_json(response.content, OcrResponse) + page_texts.append(f"--- Page {i + 1} ---\n{ocr.text.strip()}") full_text = "\n\n".join(page_texts) logger.info("Handwritten OCR extracted %d characters from %s", len(full_text), path) diff --git a/src/odoo_ocr/pipeline/mixed_unknown.py b/src/odoo_ocr/pipeline/mixed_unknown.py index ad3c44b..1dca6c9 100644 --- a/src/odoo_ocr/pipeline/mixed_unknown.py +++ b/src/odoo_ocr/pipeline/mixed_unknown.py @@ -8,19 +8,13 @@ from pathlib import Path from PIL import Image from odoo_ocr.clients import BaseVLMClient +from odoo_ocr.clients.base import load_prompt from odoo_ocr.config import Settings +from odoo_ocr.schemas import OcrResponse logger = logging.getLogger(__name__) -def _load_prompt(name: str) -> str: - from pathlib import Path - - from odoo_ocr.clients.base import BaseVLMClient - - return BaseVLMClient._load_prompt_file(Path("prompts") / name) - - async def process_mixed_unknown( path: Path, pages: list[Image.Image], @@ -32,7 +26,7 @@ async def process_mixed_unknown( For now this runs the same OCR path as scanned_print but with a more permissive prompt. Future: ensemble multiple models and merge outputs. """ - system_prompt = _load_prompt("ocr_system.txt") + system_prompt = load_prompt("ocr_system.txt") user_prompt = ( "Extract all readable text from this document. It may contain a mix of " "printed text, handwriting, stamps, or low-quality scans. Preserve layout." @@ -44,17 +38,12 @@ async def process_mixed_unknown( system_prompt=system_prompt, user_prompt=f"{user_prompt} (page {i + 1} of {len(pages)})", images=[page], + response_format=OcrResponse, temperature=0.1, model=settings.models.ocr, ) - import json - - try: - parsed = json.loads(response.content.strip().strip("`").removeprefix("json").strip()) - text = parsed.get("text", response.content) - except Exception: - text = response.content - page_texts.append(f"--- Page {i + 1} ---\n{text.strip()}") + ocr = client.parse_json(response.content, OcrResponse) + page_texts.append(f"--- Page {i + 1} ---\n{ocr.text.strip()}") full_text = "\n\n".join(page_texts) logger.info("Robust OCR extracted %d characters from %s", len(full_text), path) diff --git a/src/odoo_ocr/pipeline/scanned_print.py b/src/odoo_ocr/pipeline/scanned_print.py index 8ff34d0..e314299 100644 --- a/src/odoo_ocr/pipeline/scanned_print.py +++ b/src/odoo_ocr/pipeline/scanned_print.py @@ -8,19 +8,13 @@ from pathlib import Path from PIL import Image from odoo_ocr.clients import BaseVLMClient +from odoo_ocr.clients.base import load_prompt from odoo_ocr.config import Settings +from odoo_ocr.schemas import OcrResponse logger = logging.getLogger(__name__) -def _load_prompt(name: str) -> str: - from pathlib import Path - - from odoo_ocr.clients.base import BaseVLMClient - - return BaseVLMClient._load_prompt_file(Path("prompts") / name) - - async def process_scanned_print( path: Path, pages: list[Image.Image], @@ -28,7 +22,7 @@ async def process_scanned_print( settings: Settings, ) -> str: """Run OCR on each page of a scanned printed invoice and return combined text.""" - system_prompt = _load_prompt("ocr_system.txt") + system_prompt = load_prompt("ocr_system.txt") user_prompt = "Read all text from this printed invoice page accurately." page_texts: list[str] = [] @@ -37,18 +31,12 @@ async def process_scanned_print( system_prompt=system_prompt, user_prompt=f"{user_prompt} (page {i + 1} of {len(pages)})", images=[page], + response_format=OcrResponse, temperature=0.1, model=settings.models.ocr, ) - # The OCR prompt asks for {"text": "..."}. - import json - - try: - parsed = json.loads(response.content.strip().strip("`").removeprefix("json").strip()) - text = parsed.get("text", response.content) - except Exception: - text = response.content - page_texts.append(f"--- Page {i + 1} ---\n{text.strip()}") + ocr = client.parse_json(response.content, OcrResponse) + page_texts.append(f"--- Page {i + 1} ---\n{ocr.text.strip()}") full_text = "\n\n".join(page_texts) logger.info("OCR extracted %d characters from scanned invoice %s", len(full_text), path) diff --git a/tests/test_branches.py b/tests/test_branches.py new file mode 100644 index 0000000..bdb1bde --- /dev/null +++ b/tests/test_branches.py @@ -0,0 +1,59 @@ +"""Tests for the scanned_print, handwritten, and mixed_unknown branches.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest +from PIL import Image + +from odoo_ocr.config import Settings +from odoo_ocr.pipeline.handwritten import process_handwritten +from odoo_ocr.pipeline.mixed_unknown import process_mixed_unknown +from odoo_ocr.pipeline.scanned_print import process_scanned_print +from tests.conftest import FakeVLMClient + + +def _fake_ocr_client(text: str = "Hello invoice 123") -> FakeVLMClient: + return FakeVLMClient(responses={"OCR engine": {"text": text}}) + + +def _pages(count: int = 2) -> list[Image.Image]: + return [Image.new("RGB", (20, 20)) for _ in range(count)] + + +@pytest.mark.asyncio +async def test_scanned_print_multi_page() -> None: + settings = Settings() + client = _fake_ocr_client() + text = await process_scanned_print(Path("scan.jpg"), _pages(2), client, settings) + assert "--- Page 1 ---" in text + assert "--- Page 2 ---" in text + assert text.count("Hello invoice 123") == 2 + assert all(call["model"] == settings.models.ocr for call in client.calls) + + +@pytest.mark.asyncio +async def test_scanned_print_rejects_invalid_json() -> None: + """Branches no longer silently fall back to raw content; invalid JSON fails.""" + client = FakeVLMClient(responses={"OCR engine": "plain text without json"}) + with pytest.raises(ValueError, match="Response does not match schema"): + await process_scanned_print(Path("scan.jpg"), _pages(1), client, Settings()) + + +@pytest.mark.asyncio +async def test_handwritten_extracts_text() -> None: + settings = Settings() + client = _fake_ocr_client("handwritten total 42") + text = await process_handwritten(Path("hand.jpg"), _pages(1), client, settings) + assert "handwritten total 42" in text + assert client.calls[0]["model"] == settings.models.ocr + + +@pytest.mark.asyncio +async def test_mixed_unknown_extracts_text() -> None: + settings = Settings() + client = _fake_ocr_client("mixed content text") + text = await process_mixed_unknown(Path("mixed.jpg"), _pages(1), client, settings) + assert "mixed content text" in text + assert client.calls[0]["model"] == settings.models.ocr