mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-19 20:17:12 +02:00
Victor Kuznetsov
e938b57f8c
Move the draft-annotation logic (PaddleOCR line detection, word-box grouping, three script-chosen recognition engines, crop-jitter stability gate) from the evaluation-only scripts into the installable package, with lazy paddle imports and a text-draft extra (CPU, no torch). draft_text_lines() returns accepted (crop-stable, NEVER ground-truth-correct - precision on the reference posters was 90.0% / 94.4%) and rejected proposals; source_pixel_sha256 is re-exported for manifest building. scripts/infer_text_lines.py now dogfoods the package module instead of loading the eval script by path.
100 lines
3.5 KiB
Python
100 lines
3.5 KiB
Python
# /// script
|
|
# requires-python = ">=3.11"
|
|
# dependencies = [
|
|
# "click",
|
|
# "huggingface-hub>=0.20.0",
|
|
# "numpy",
|
|
# "onnxruntime>=1.24.0",
|
|
# "opencv-python-headless<5",
|
|
# "paddleocr>=3.3.3",
|
|
# "paddlepaddle",
|
|
# "pillow",
|
|
# ]
|
|
# ///
|
|
"""Infer stable source-text lines without modifying an image.
|
|
|
|
This evaluation-only dry run proposes line annotations for selective text
|
|
restoration. Every proposal still needs human verification: stable OCR can lose
|
|
punctuation with high confidence. It separately flags lines whose recognition
|
|
changes under crop jitter or whose minimum confidence is below the threshold.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import click
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
log = logging.getLogger(__name__)
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
sys.path.insert(0, str(ROOT))
|
|
|
|
# Dogfoods the packaged draft API (remove_ai_watermarks.text_draft); this script
|
|
# keeps only the CLI wrapper so the package stays the one home for the logic.
|
|
from remove_ai_watermarks.text_draft import ( # noqa: E402
|
|
_detect_line_boxes,
|
|
_recognize,
|
|
choose_language,
|
|
stable_recognition,
|
|
)
|
|
|
|
|
|
@click.command()
|
|
@click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path))
|
|
@click.option("--out", required=True, type=click.Path(dir_okay=False, path_type=Path))
|
|
@click.option("--min-score", default=0.85, show_default=True, type=click.FloatRange(0.0, 1.0))
|
|
def main(source: Path, out: Path, min_score: float) -> None:
|
|
"""Write draft line text for SOURCE; manually verify every proposal."""
|
|
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
|
|
os.environ["PADDLE_PDX_DISABLE_MODEL_SOURCE_CHECK"] = "True"
|
|
from paddleocr import PaddleOCR, TextRecognition
|
|
|
|
source_rgb = np.asarray(Image.open(source).convert("RGB"))
|
|
detector = PaddleOCR(
|
|
lang="ch",
|
|
use_doc_orientation_classify=False,
|
|
use_doc_unwarping=False,
|
|
use_textline_orientation=False,
|
|
)
|
|
engines = {
|
|
"en": TextRecognition(model_name="en_PP-OCRv5_mobile_rec"),
|
|
"ru": TextRecognition(model_name="eslav_PP-OCRv5_mobile_rec"),
|
|
"ch": TextRecognition(model_name="PP-OCRv5_server_rec"),
|
|
}
|
|
boxes = _detect_line_boxes(detector, source_rgb)
|
|
accepted = []
|
|
rejected = []
|
|
for box in boxes:
|
|
probes = {}
|
|
for language, engine in engines.items():
|
|
script = "cjk" if language == "ch" else "alphabetic"
|
|
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 (0.08, 0.12, 0.2)]
|
|
text = stable_recognition(reads, min_score)
|
|
result = {
|
|
"box": box,
|
|
"script": script,
|
|
"language": language,
|
|
"reads": [{"text": value, "score": score} for value, score in reads],
|
|
}
|
|
if text is None:
|
|
rejected.append(result)
|
|
else:
|
|
accepted.append({"box": box, "text": text, "script": script, "min_score": min(score for _, score in reads)})
|
|
payload = {"source": source.name, "accepted": accepted, "rejected": rejected}
|
|
out.parent.mkdir(parents=True, exist_ok=True)
|
|
out.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
log.info("Accepted %s lines and rejected %s uncertain lines", len(accepted), len(rejected))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|