ocr-foundry 0.2.0__py3-none-any.whl

khocr_gen/lmdb_pack.py

@@ -0,0 +1,132 @@
+"""Pack a labels.txt + image directory into an LMDB database.
+
+Key format:
+    b"num-samples"       -> ASCII count of samples written
+    b"image-{i:09d}"     -> raw JPEG bytes for sample *i* (1-indexed)
+    b"label-{i:09d}"     -> UTF-8 text label for sample *i* (1-indexed)
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+from tqdm import tqdm
+
+_DEFAULT_COMMIT_EVERY_BYTES = 256 * 1024 * 1024  # 256 MiB
+
+
+def _count_non_empty_lines(path: Path) -> int:
+    count = 0
+    with path.open("r", encoding="utf-8") as fh:
+        for line in fh:
+            if line.strip():
+                count += 1
+    return count
+
+
+def pack_lmdb(
+    labels_file: str | Path,
+    images_dir: str | Path,
+    out_dir: str | Path,
+    *,
+    jpeg_quality: int = 90,
+    map_size_gb: int = 256,
+    verbose: bool = False,
+    commit_every: int = 5_000,
+    commit_every_bytes: int = _DEFAULT_COMMIT_EVERY_BYTES,
+) -> int:
+    """Write an LMDB database from *labels_file* and *images_dir*.
+
+    Pending encoded images are flushed to LMDB when either *commit_every*
+    samples or *commit_every_bytes* of encoded image data have accumulated,
+    whichever comes first -- a count-only bound lets a run of large images
+    balloon the transient buffer to gigabytes before a single flush.
+
+    Returns the number of samples written.
+    """
+    import cv2
+    import lmdb
+
+    labels_path = Path(labels_file).expanduser().resolve()
+    images_path = Path(images_dir).expanduser().resolve()
+    out_path = Path(out_dir).expanduser().resolve()
+
+    if not labels_path.is_file():
+        raise FileNotFoundError(f"Labels file not found: {labels_path}")
+    if not images_path.is_dir():
+        raise FileNotFoundError(f"Images directory not found: {images_path}")
+
+    out_path.mkdir(parents=True, exist_ok=True)
+
+    map_size = map_size_gb * (1024**3)
+    env = lmdb.open(str(out_path), map_size=map_size)
+
+    samples_written = 0
+    errors = 0
+    total_lines = _count_non_empty_lines(labels_path)
+
+    # Batch writes to avoid holding the LMDB write lock for too long
+    pending_pairs: list[tuple[bytes, bytes, bytes, bytes]] = []
+    pending_bytes = 0
+
+    def _flush(txn, pairs: list) -> None:
+        for img_key, img_bytes, lbl_key, lbl_bytes in pairs:
+            txn.put(img_key, img_bytes)
+            txn.put(lbl_key, lbl_bytes)
+        pairs.clear()
+
+    with labels_path.open("r", encoding="utf-8") as fh:
+        lines = (ln.rstrip("\n") for ln in fh if ln.strip())
+
+        for line in tqdm(lines, total=total_lines, desc="  Packing LMDB", unit="img"):
+            parts = line.split("\t", 1)
+            if len(parts) < 2:
+                continue
+            img_name, label = parts[0].strip(), parts[1]
+            img_path = images_path / img_name
+            if not img_path.exists():
+                # Try next to the labels file
+                img_path = labels_path.parent / img_name
+            if not img_path.exists():
+                if verbose:
+                    print(f"  [SKIP] image not found: {img_name}", file=sys.stderr)
+                errors += 1
+                continue
+
+            img = cv2.imread(str(img_path), cv2.IMREAD_UNCHANGED)
+            if img is None:
+                if verbose:
+                    print(f"  [SKIP] cv2.imread failed: {img_name}", file=sys.stderr)
+                errors += 1
+                continue
+
+            ok, buf = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, jpeg_quality])
+            if not ok:
+                errors += 1
+                continue
+
+            samples_written += 1
+            idx = samples_written
+            img_key = f"image-{idx:09d}".encode()
+            lbl_key = f"label-{idx:09d}".encode()
+            img_bytes = bytes(buf)
+            pending_pairs.append((img_key, img_bytes, lbl_key, label.encode("utf-8")))
+            pending_bytes += len(img_bytes)
+
+            if len(pending_pairs) >= commit_every or pending_bytes >= commit_every_bytes:
+                with env.begin(write=True) as txn:
+                    _flush(txn, pending_pairs)
+                pending_bytes = 0
+
+    # Flush remaining and write final count
+    with env.begin(write=True) as txn:
+        _flush(txn, pending_pairs)
+        txn.put(b"num-samples", str(samples_written).encode())
+
+    env.close()
+
+    if verbose and errors:
+        print(f"  Skipped {errors} samples due to missing/corrupt images", file=sys.stderr)
+
+    return samples_written

khocr_gen/logging.py

@@ -0,0 +1,25 @@
+"""Logging utilities for khocr-gen."""
+
+from __future__ import annotations
+
+import logging
+
+_DEFAULT_LOG_FORMAT = "%(asctime)s | %(levelname)s | %(name)s | %(message)s"
+
+
+def configure_logging(level: int = logging.INFO) -> None:
+    """Configure package logging once using a consistent format."""
+    root_logger = logging.getLogger("khocr_gen")
+    if root_logger.handlers:
+        root_logger.setLevel(level)
+        return
+
+    handler = logging.StreamHandler()
+    handler.setFormatter(logging.Formatter(_DEFAULT_LOG_FORMAT))
+    root_logger.addHandler(handler)
+    root_logger.setLevel(level)
+
+
+def get_logger(name: str) -> logging.Logger:
+    """Return a child logger rooted at khocr_gen."""
+    return logging.getLogger(f"khocr_gen.{name}")

khocr_gen/normalizer.py

@@ -0,0 +1,221 @@
+"""Khmer text normalization via khmernormalizer.
+
+Wraps https://github.com/seanghay/khmernormalizer with configurable parameters exposed to the CLI and config file,
+so users can tune normalization for their corpus.
+"""
+
+from __future__ import annotations
+
+import unicodedata
+from dataclasses import asdict, dataclass, fields
+from typing import Any
+
+_HAS_KHMER_NORMALIZER = False
+_normalize_fn: Any = None
+_TextFixerConfig: Any = None
+
+
+def _ensure_backend() -> tuple[Any, Any]:
+    global _HAS_KHMER_NORMALIZER, _normalize_fn, _TextFixerConfig
+    if _HAS_KHMER_NORMALIZER:
+        return _normalize_fn, _TextFixerConfig
+
+    try:
+        from khmernormalizer import TextFixerConfig as _TFC
+        from khmernormalizer import normalize as _nf
+
+        _normalize_fn = _nf
+        _TextFixerConfig = _TFC
+        _HAS_KHMER_NORMALIZER = True
+    except ImportError as exc:
+        raise ImportError(
+            "khmernormalizer is required for text normalization. "
+            "Install with: pip install khmernormalizer"
+        ) from exc
+    return _normalize_fn, _TextFixerConfig
+
+
+@dataclass
+class NormalizerConfig:
+    """Parameters for Khmer text normalization.
+
+    All parameters map directly to `khmernormalizer.normalize` kwargs.
+    """
+
+    policy: str = "legacy-khmer"  # legacy-khmer | nfc | none
+    unicode_norm: str = ""
+    emoji_replacement: str | None = ""
+    url_replacement: str | None = ""
+    remove_zwsp: bool = True
+
+    # ftfy TextFixerConfig fields
+    unescape_html: str = "auto"
+    remove_terminal_escapes: bool = True
+    fix_encoding: bool = True
+    restore_byte_a0: bool = True
+    replace_lossy_sequences: bool = True
+    decode_inconsistent_utf8: bool = True
+    fix_c1_controls: bool = True
+    fix_latin_ligatures: bool = True
+    fix_character_width: bool = True
+    uncurl_quotes: bool = True
+    fix_line_breaks: bool = True
+    fix_surrogates: bool = True
+    remove_control_chars: bool = True
+    normalization: str = "NFC"
+    max_decode_length: int = 1_000_000
+    explain: bool = False
+
+    # Shorthand for "apply no normalization at all"
+    passthrough: bool = False
+
+    def __post_init__(self) -> None:
+        if self.policy not in {"legacy-khmer", "nfc", "none"}:
+            raise ValueError("normalization policy must be legacy-khmer, nfc, or none")
+
+    def to_fix_text_config(self) -> Any:
+        """Build a `TextFixerConfig` from this config."""
+        _, TFC = _ensure_backend()
+        return TFC(
+            unescape_html=self.unescape_html,
+            remove_terminal_escapes=self.remove_terminal_escapes,
+            fix_encoding=self.fix_encoding,
+            restore_byte_a0=self.restore_byte_a0,
+            replace_lossy_sequences=self.replace_lossy_sequences,
+            decode_inconsistent_utf8=self.decode_inconsistent_utf8,
+            fix_c1_controls=self.fix_c1_controls,
+            fix_latin_ligatures=self.fix_latin_ligatures,
+            fix_character_width=self.fix_character_width,
+            uncurl_quotes=self.uncurl_quotes,
+            fix_line_breaks=self.fix_line_breaks,
+            fix_surrogates=self.fix_surrogates,
+            remove_control_chars=self.remove_control_chars,
+            normalization=self.normalization,
+            max_decode_length=self.max_decode_length,
+            explain=self.explain,
+        )
+
+    @staticmethod
+    def add_args(parser: Any, *, prefix: str = "--norm-") -> None:
+        """Register normalizer flags on an argparse parser.
+
+        Args:
+            parser: An `argparse.ArgumentParser` or argument group.
+            prefix: Prefix for flag names (default ``--norm-``).
+        """
+        g = parser.add_argument_group("Text normalization")
+        g.add_argument(
+            f"{prefix}policy",
+            choices=["legacy-khmer", "nfc", "none"],
+            default=None,
+            help="Normalization policy (default: legacy-khmer, or nfc with --multilingual)",
+        )
+        g.add_argument(
+            f"{prefix}unicode-norm",
+            default="",
+            choices=["NFKC", "NFC", "NFD", "NFKD", "none"],
+            help="Unicode normalization form (default: none). khmernormalizer already normalizes internally.",
+        )
+        g.add_argument(
+            f"{prefix}emoji-replacement",
+            default="",
+            metavar="STR",
+            help="Replacement string for emoji characters (default: '' = remove)",
+        )
+        g.add_argument(
+            f"{prefix}url-replacement",
+            default="",
+            metavar="STR",
+            help="Replacement string for URLs (default: '' = remove)",
+        )
+        g.add_argument(
+            f"{prefix}no-remove-zwsp",
+            action="store_true",
+            help="Keep zero-width spaces instead of removing them",
+        )
+        g.add_argument(
+            f"{prefix}no-fix-encoding",
+            action="store_true",
+            help="Disable ftfy encoding fixes",
+        )
+        g.add_argument(
+            f"{prefix}no-uncurl-quotes",
+            action="store_true",
+            help="Disable quote uncurling",
+        )
+        g.add_argument(
+            f"{prefix}no-fix-line-breaks",
+            action="store_true",
+            help="Disable line break normalization",
+        )
+        g.add_argument(
+            f"{prefix}passthrough",
+            action="store_true",
+            help="Skip all normalization (text passes through unchanged)",
+        )
+
+    @classmethod
+    def from_args(cls, args: Any, *, prefix: str = "norm_") -> NormalizerConfig:
+        """Build a `NormalizerConfig` from parsed CLI args."""
+
+        def _get(name: str) -> Any:
+            return getattr(args, f"{prefix}{name}", None)
+
+        unicode_norm_val = _get("unicode_norm")
+        if unicode_norm_val and str(unicode_norm_val).lower() == "none":
+            unicode_norm_val = ""
+
+        return cls(
+            policy=_get("policy")
+            or ("nfc" if getattr(args, "multilingual", False) else "legacy-khmer"),
+            unicode_norm=unicode_norm_val if unicode_norm_val else "",
+            emoji_replacement=_get("emoji_replacement"),
+            url_replacement=_get("url_replacement"),
+            remove_zwsp=not bool(_get("no_remove_zwsp")),
+            fix_encoding=not bool(_get("no_fix_encoding")),
+            uncurl_quotes=not bool(_get("no_uncurl_quotes")),
+            fix_line_breaks=not bool(_get("no_fix_line_breaks")),
+            passthrough=bool(_get("passthrough")),
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialise to a plain dict for IPC."""
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]) -> NormalizerConfig:
+        """Reconstruct from a dict."""
+        known = {f.name for f in fields(cls)}
+        return cls(**{key: value for key, value in d.items() if key in known})
+
+
+def normalize(text: str, cfg: NormalizerConfig | None = None) -> str:
+    """Normalize Khmer text.
+
+    Args:
+        text: Raw input text.
+        cfg: Normalizer config. If *None*, a default `NormalizerConfig` is used.
+
+    Returns:
+        Normalized text, or the original text if `cfg.passthrough` is `True`.
+    """
+    if cfg is None:
+        cfg = NormalizerConfig()
+
+    if cfg.passthrough or cfg.policy == "none":
+        return text
+    if cfg.policy == "nfc":
+        return unicodedata.normalize("NFC", text)
+
+    nf, _ = _ensure_backend()
+
+    kwargs: dict[str, Any] = {}
+    if cfg.unicode_norm:
+        kwargs["unicode_norm"] = cfg.unicode_norm
+
+    kwargs["emoji_replacement"] = cfg.emoji_replacement
+    kwargs["url_replacement"] = cfg.url_replacement
+    kwargs["remove_zwsp"] = cfg.remove_zwsp
+    kwargs["fix_text_config"] = cfg.to_fix_text_config()
+
+    return nf(text, **kwargs)