"""text_draft: proposal-only OCR for verified-text manifests (no paddle needed).""" from __future__ import annotations from typing import TYPE_CHECKING, Any from PIL import Image from remove_ai_watermarks.text_draft import ( TextDraft, choose_language, draft_available, draft_text_lines, group_word_boxes, stable_recognition, ) if TYPE_CHECKING: from pathlib import Path class TestPureHelpers: def test_choose_language_follows_unicode_script(self): assert choose_language({"ch": ("你好", 0.9), "ru": ("xxx", 0.1), "en": ("yyy", 0.1)}) == "ch" assert choose_language({"ch": ("??", 0.1), "ru": ("привет", 0.9), "en": ("yyy", 0.1)}) == "ru" assert choose_language({"ch": ("??", 0.1), "ru": ("xxx", 0.1), "en": ("hello", 0.9)}) == "en" def test_stable_recognition_requires_identical_normalizations(self): # Same text modulo spacing/case (comma kept): stable. assert ( stable_recognition([("Hello, World", 0.97), ("hello,world", 0.96), ("HELLO, WORLD", 0.95)]) == "Hello, World" ) # Different texts: rejected (None). assert stable_recognition([("Hello", 0.97), ("Hallo", 0.96), ("Hello", 0.95)]) is None # Identical text but one read under the floor: rejected. assert stable_recognition([("Hello", 0.97), ("Hello", 0.60), ("Hello", 0.95)]) is None assert stable_recognition([("Hello", 0.97), ("Hello", 0.96), ("Hello", 0.95)], min_score=0.99) is None def test_group_word_boxes_merges_a_line_and_keeps_rows_apart(self): # Two words on one baseline merge; a distant lower line stays separate. merged = group_word_boxes([(10, 100, 60, 140), (70, 100, 120, 140)]) assert len(merged) == 1 assert merged[0] == (10, 100, 120, 140) apart = group_word_boxes([(10, 100, 60, 140), (10, 300, 60, 340)]) assert len(apart) == 2 class TestDraftAvailable: def test_returns_a_bool_without_importing_paddle(self): assert isinstance(draft_available(), bool) class _FakeDetector: """Emits Paddle-style pages for one synthetic poster.""" def __init__(self, boxes: list[tuple[int, int, int, int]]) -> None: self._boxes = boxes def predict(self, _image: Any) -> Any: import numpy as np yield { "rec_boxes": np.asarray(self._boxes, dtype=np.float32), "rec_scores": [0.99] * len(self._boxes), } class _FakeEngine: """Returns a fixed text at high score regardless of the crop.""" def __init__(self, text: str) -> None: self._text = text def predict(self, _crop: Any) -> Any: yield {"rec_text": self._text, "rec_score": 0.97} class _JitterEngine: """Changes its answer with the crop height - exactly what the gate rejects.""" def __init__(self, texts: list[str]) -> None: self._texts = texts self._calls = 0 def predict(self, crop: Any) -> Any: text = self._texts[self._calls % len(self._texts)] self._calls += 1 yield {"rec_text": text, "rec_score": 0.97} class TestDraftTextLines: @staticmethod def _poster(tmp_path: Path) -> Path: path = tmp_path / "poster.png" Image.new("RGB", (400, 300), (255, 255, 255)).save(path) return path def test_accepts_crop_stable_and_rejects_jittered(self, tmp_path: Path): path = self._poster(tmp_path) stable = _FakeEngine("Invoice 42") jitter = _JitterEngine(["Hello", "Hallo", "Hullo"]) # All three probe engines read the same string, so language = en; the # jitter engine then serves as the en engine and flips across crops. draft = draft_text_lines( path, detector=_FakeDetector([(20, 20, 200, 60), (20, 80, 200, 120)]), engines={"en": jitter, "ru": stable, "ch": stable}, ) assert isinstance(draft, TextDraft) assert len(draft.rejected) == 2 assert all(line.language == "en" for line in draft.rejected) assert draft.accepted == () def test_accepted_line_carries_box_text_script_and_floor(self, tmp_path: Path): path = self._poster(tmp_path) engine = _FakeEngine("Total: 1,234.56") draft = draft_text_lines( path, detector=_FakeDetector([(20, 20, 260, 60)]), engines={"en": engine, "ru": engine, "ch": engine}, ) (line,) = draft.accepted assert line.box == (20, 20, 260, 60) assert line.text == "Total: 1,234.56" assert line.script == "alphabetic" assert line.language == "en" assert line.min_score >= 0.85 def test_cjk_probe_switches_script_and_language(self, tmp_path: Path): path = self._poster(tmp_path) cjk = _FakeEngine("每天都是一个新的机会。") latin = _FakeEngine("hello") draft = draft_text_lines( path, detector=_FakeDetector([(30, 30, 300, 90)]), engines={"en": latin, "ru": latin, "ch": cjk}, ) (line,) = draft.accepted assert line.language == "ch" assert line.script == "cjk"