diff --git a/src/remove_ai_watermarks/_internal/text_restoration.py b/src/remove_ai_watermarks/_internal/text_restoration.py index b956add..58431d2 100644 --- a/src/remove_ai_watermarks/_internal/text_restoration.py +++ b/src/remove_ai_watermarks/_internal/text_restoration.py @@ -185,6 +185,20 @@ def restore_verified_text( return Image.fromarray(restored) +def _glyph_crop_box(box: tuple[int, int, int, int], width: int, height: int) -> tuple[int, int, int, int]: + """Widen the detector box so descenders stay inside the silhouette crop. + + Paddle line boxes sit 2-5 px above the true ink bottom on the poster + fixtures. The crop is the box itself, so those pixels never reached the + donor composite. Expand only on Y: 8% up, 25% down, clamped to the frame. + """ + x1, y1, x2, y2 = box + line_h = max(1, y2 - y1) + pad_top = max(1, round(line_h * 0.08)) + pad_bot = max(2, round(line_h * 0.25)) + return max(0, x1), max(0, y1 - pad_top), min(width, x2), min(height, y2 + pad_bot) + + def source_silhouette_mask( source_rgb: NDArray[Any], box: tuple[int, int, int, int], @@ -192,7 +206,7 @@ def source_silhouette_mask( ) -> NDArray[Any]: """Recover a thresholded glyph shape without retaining source amplitudes.""" height, width = source_rgb.shape[:2] - x1, y1, x2, y2 = _clip_box(box, width, height) + x1, y1, x2, y2 = _glyph_crop_box(box, width, height) gray = cv2.cvtColor(source_rgb[y1:y2, x1:x2], cv2.COLOR_RGB2GRAY) support = np.ones(gray.shape, dtype=np.uint8) if angle: diff --git a/src/remove_ai_watermarks/text_draft.py b/src/remove_ai_watermarks/text_draft.py index 96256a6..742c009 100644 --- a/src/remove_ai_watermarks/text_draft.py +++ b/src/remove_ai_watermarks/text_draft.py @@ -223,15 +223,22 @@ def draft_text_lines( min_score: float = 0.85, detector: Any | None = None, engines: dict[str, Any] | None = None, + stable: bool = True, ) -> TextDraft: """Propose verified-text manifest lines for ``image``; never verified ones. Args: image: path of the source image (draft boxes are in ITS pixel space). - min_score: recognition confidence floor for every jittered read. + min_score: recognition confidence floor. With ``stable=True`` every + jittered read must clear it; with ``stable=False`` the single + probe for the chosen language must. detector/engines: injectable Paddle objects (tests use fakes); when None they are built from the ``text-draft`` extra, which must be installed (``draft_available()`` reports it). + stable: when True (default), accept a line only if three crop paddings + normalize identically. When False, take one recognition trio and + accept on score alone. Automatic restoration uses box/script only, + so the jitter gate is optional there. Returns: ``TextDraft`` with crop-stable ``accepted`` proposals and unstable @@ -262,8 +269,14 @@ def draft_text_lines( probes[language] = _recognize(engine, source_rgb, box, script, 0.1) language = choose_language(probes) script = "cjk" if language == "ch" else "alphabetic" - reads = [_recognize(engines[language], source_rgb, box, script, ratio) for ratio in JITTER_RATIOS] - text = stable_recognition(reads, min_score) + if stable: + reads = [_recognize(engines[language], source_rgb, box, script, ratio) for ratio in JITTER_RATIOS] + text = stable_recognition(reads, min_score) + else: + text, score = probes[language] + reads = [(text, score)] + if not text.strip() or score < min_score: + text = None if text is None: rejected.append( RejectedLine( diff --git a/tests/test_text_draft.py b/tests/test_text_draft.py index 6a1abc0..4a8f3aa 100644 --- a/tests/test_text_draft.py +++ b/tests/test_text_draft.py @@ -139,3 +139,20 @@ class TestDraftTextLines: (line,) = draft.accepted assert line.language == "ch" assert line.script == "cjk" + + def test_unstable_geometry_draft_accepts_a_single_high_score_read(self, tmp_path: Path): + path = self._poster(tmp_path) + jitter = _JitterEngine(["Hello", "Hallo", "Hullo"]) + stable = _FakeEngine("Hello") + # Probes disagree across engines so language is en; jitter would fail + # the stable gate. Geometry mode uses the one en probe and accepts. + draft = draft_text_lines( + path, + min_score=0.75, + stable=False, + detector=_FakeDetector([(20, 20, 200, 60)]), + engines={"en": jitter, "ru": stable, "ch": stable}, + ) + (line,) = draft.accepted + assert line.text == "Hello" + assert line.min_score >= 0.75 diff --git a/tests/test_text_restoration.py b/tests/test_text_restoration.py index 9ea4887..bdc0e2e 100644 --- a/tests/test_text_restoration.py +++ b/tests/test_text_restoration.py @@ -15,6 +15,7 @@ from remove_ai_watermarks._internal.text_restoration import ( load_verified_text_manifest, restore_verified_text, source_pixel_sha256, + source_silhouette_mask, ) @@ -126,3 +127,14 @@ def test_restoration_uses_lama_and_qwen_vae_core(monkeypatch) -> None: assert calls assert np.all(restored[16, 20] == donor[16, 20]) assert np.all(restored[0, 0] == candidate[0, 0]) + + +def test_silhouette_includes_descender_below_the_detector_box() -> None: + source = np.full((40, 50, 3), 240, dtype=np.uint8) + source[10:22, 18:24] = 20 + source[22:27, 18:22] = 20 # tail of a y / Cyrillic u, under the box + + mask = source_silhouette_mask(source, (10, 10, 40, 22)) + + assert mask[24, 20] == 255 + assert mask[16, 20] == 255