anaksor-ocr 0.1.0__py3-none-any.whl

anaksor/__init__.py

@@ -0,0 +1,54 @@
+"""anaksor: Khmer OCR (អានអក្សរ, "read the letters").
+
+import anaksor
+
+ocr = anaksor.OCR()                          # weights download on first use
+result = ocr("page.jpg")
+print(result["text"])
+
+pages = ocr.read("scan.pdf", keep_images=True)            # needs anaksor-ocr[pdf]
+open("scan.pdf", "wb").write(anaksor.export.to_pdf(pages))   # searchable PDF
+"""
+
+from __future__ import annotations
+
+import threading
+from importlib.metadata import PackageNotFoundError, version
+
+from . import export
+from .detector import Detector
+from .inputs import InputError, load_image
+from .models import WeightsError
+from .ocr import OCR
+from .recognizer import Recognizer
+from .text import normalize
+
+try:
+    __version__ = version("anaksor-ocr")
+except PackageNotFoundError:  # running from a source tree
+    __version__ = "0.0.0"
+
+__all__ = [
+    "OCR",
+    "Detector",
+    "InputError",
+    "Recognizer",
+    "WeightsError",
+    "__version__",
+    "export",
+    "load_image",
+    "normalize",
+    "read",
+]
+
+_default: OCR | None = None
+_default_lock = threading.Lock()
+
+
+def read(src, **kwargs) -> list[dict]:
+    """OCR an image or PDF with a shared default ``OCR()``; returns one result per page."""
+    global _default
+    with _default_lock:
+        if _default is None:
+            _default = OCR()
+    return _default.read(src, **kwargs)

anaksor/__main__.py

@@ -0,0 +1,3 @@
+from .cli import main
+
+raise SystemExit(main())

anaksor/cli.py

@@ -0,0 +1,151 @@
+"""Command line: ``anaksor scan.pdf -o scan.docx``."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+FORMATS = ("text", "json", "pdf", "docx")
+_EXT = {"text": ".txt", "json": ".json", "pdf": ".pdf", "docx": ".docx"}
+_BY_EXT = {".txt": "text", ".md": "text", ".json": "json", ".pdf": "pdf", ".docx": "docx"}
+
+
+def parse_pages(spec: str | None) -> list[int] | None:
+    """'1-3,5' to [1, 2, 3, 5]."""
+    if not spec:
+        return None
+    pages: list[int] = []
+    for part in spec.split(","):
+        part = part.strip()
+        if not part:
+            continue
+        a, sep, b = part.partition("-")
+        try:
+            lo, hi = int(a), int(b) if sep else int(a)
+        except ValueError:
+            raise argparse.ArgumentTypeError(f"bad page range {part!r} (use e.g. 1-3,5)") from None
+        if lo < 1 or hi < lo:
+            raise argparse.ArgumentTypeError(f"bad page range {part!r}")
+        pages.extend(range(lo, hi + 1))
+    return pages
+
+
+def build_parser() -> argparse.ArgumentParser:
+    from . import __version__
+
+    p = argparse.ArgumentParser(
+        prog="anaksor",
+        description="Khmer OCR: images and PDFs to text, JSON, searchable PDF or Word.",
+        epilog="Examples:\n  anaksor photo.jpg\n  anaksor scan.pdf -o scan.docx\n  anaksor a.jpg b.png -o out/ -f json",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    p.add_argument("inputs", nargs="*", type=Path, help="images (JPEG, PNG, WEBP, BMP, TIFF) or PDFs")
+    p.add_argument(
+        "-o", "--output", type=Path, help="output file, or a directory for several inputs (default: stdout for text/json)"
+    )
+    p.add_argument("-f", "--format", choices=FORMATS, help="default: from the output's extension, else text")
+    p.add_argument("--pages", type=parse_pages, help="PDF pages to read, e.g. 1-3,5 (default: all)")
+    p.add_argument("--dpi", type=int, default=200, help="PDF render resolution (default: 200)")
+    p.add_argument("--device", default="auto", help="auto, cpu, cuda, cuda:1, mps (default: auto)")
+    p.add_argument(
+        "--backend", choices=("auto", "torch", "onnx"), default="auto", help="torch or onnx runtime (default: torch if installed)"
+    )
+    p.add_argument("--det-weights", type=Path, help="detection weights file (default: cached download)")
+    p.add_argument("--rec-weights", type=Path, help="recognition weights file (default: cached download)")
+    p.add_argument("--pad-y", type=float, help="line crop margin above and below, fraction of line height")
+    p.add_argument("--pad-x", type=float, help="line crop margin left and right")
+    p.add_argument("--min-score", type=float, default=0.5, help="drop lines read below this confidence")
+    p.add_argument("--highlight-low", action="store_true", help="Word: highlight low-confidence lines")
+    p.add_argument("--fetch-models", action="store_true", help="download the model weights and exit")
+    p.add_argument("--version", action="version", version=f"anaksor {__version__}")
+    return p
+
+
+def _out_path(args, src: Path, fmt: str, many: bool) -> Path | None:
+    if args.output is None:
+        return None
+    if many or args.output.is_dir():
+        args.output.mkdir(parents=True, exist_ok=True)
+        return args.output / (src.stem + _EXT[fmt])
+    return args.output
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+
+    from .models import DET_DEFAULT, REC_DEFAULT, WeightsError, load_spec, resolve_weights
+
+    if args.fetch_models:
+        try:
+            from .runtime import choose_backend
+
+            for name, w in ((DET_DEFAULT, args.det_weights), (REC_DEFAULT, args.rec_weights)):
+                fmt = choose_backend(args.backend, w)
+                print(f"{name}: {resolve_weights(load_spec(name), w, fmt=fmt)}")
+        except WeightsError as e:
+            print(f"anaksor: {e}", file=sys.stderr)
+            return 1
+        return 0
+    if not args.inputs:
+        parser.error("give at least one image or PDF")
+    fmt = args.format or _BY_EXT.get(args.output.suffix.lower() if args.output else "", "text")
+    many = len(args.inputs) > 1
+    if fmt in ("pdf", "docx") and args.output is None:
+        parser.error(f"{fmt} output needs -o FILE")
+    if many and args.output is not None and args.output.suffix and not args.output.is_dir():
+        parser.error("several inputs: -o must be a directory")
+    for src in args.inputs:
+        if not src.is_file():
+            parser.error(f"no such file: {src}")
+
+    from . import export
+    from .inputs import InputError
+    from .ocr import OCR
+
+    try:
+        ocr = OCR(
+            device=args.device,
+            backend=args.backend,
+            det_weights=args.det_weights,
+            rec_weights=args.rec_weights,
+            pad_y=args.pad_y,
+            pad_x=args.pad_x,
+            min_rec_score=args.min_score,
+        )
+    except (WeightsError, ValueError, ImportError) as e:
+        print(f"anaksor: {e}", file=sys.stderr)
+        return 1
+
+    status = 0
+    for src in args.inputs:
+        try:
+            pages = ocr.read(src, pages=args.pages, dpi=args.dpi, keep_images=fmt == "pdf")
+        except (InputError, ImportError) as e:
+            print(f"anaksor: {src}: {e}", file=sys.stderr)
+            status = 1
+            continue
+        out = _out_path(args, src, fmt, many)
+        if fmt == "text":
+            data: str | bytes = export.to_text(pages) + "\n"
+        elif fmt == "json":
+            data = json.dumps({"file": str(src), "pages": pages}, ensure_ascii=False, indent=1) + "\n"
+        elif fmt == "pdf":
+            data = export.to_pdf(pages, title=src.stem)
+        else:
+            data = export.to_docx(pages, title=src.stem, highlight_low=args.highlight_low)
+        if out is None:
+            sys.stdout.write(data if isinstance(data, str) else data.decode())
+        elif isinstance(data, bytes):
+            out.write_bytes(data)
+        else:
+            out.write_text(data, encoding="utf-8")
+        if out is not None:
+            print(f"{src} -> {out}", file=sys.stderr)
+    return status
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

anaksor/data/det_repvit_db_v2.yml

@@ -0,0 +1,33 @@
+# Khmer text detection v2: RepViT backbone + DBNet (KhmerST + synthetic pages).
+# Architecture block as in Alpha-OCR configs/repvit_db_khmer_v2.yml (training-only keys are ignored).
+name: det-repvit-db-v2
+kind: det
+weights:
+  file: khmer_det_repvit_db_v2.safetensors
+  sha256: null            # set when the release weights are published
+  onnx: khmer_det_repvit_db_v2.onnx   # tools/export_onnx.py; the default (torch-free) install runs this
+  onnx_sha256: null
+input:
+  limit_side_len: 1280    # long side resized down to this (small document text needs 1280+)
+  min_side_len: 736       # smaller images are upscaled
+  mean: [0.485, 0.456, 0.406]
+  std: [0.229, 0.224, 0.225]
+architecture:
+  algorithm: DB_mobile
+  Backbone:
+    name: RepSVTR_det
+  Neck:
+    name: RSEFPN
+    out_channels: 96
+    shortcut: true
+  Head:
+    name: DBHead
+    k: 50
+postprocess:
+  name: DBPostProcess
+  thresh: 0.3
+  box_thresh: 0.5
+  max_candidates: 1000
+  unclip_ratio: 1.5
+  score_mode: slow
+  use_dilation: false     # dilating at inference glues tightly spaced Khmer lines together

anaksor/data/det_repvit_db_v3.yml

@@ -0,0 +1,22 @@
+# Khmer text detection v3: RepViT backbone + DBNet (KhmerST, synthetic pages and signs,
+# HierText English on photos). Made with tools/spec_from_config.py from the training config
+# published as det_latest.yml at Sokheng/anaksor-ocr, tag rec-v4-det-v3.
+name: det-repvit-db-v3
+kind: det
+weights:
+  file: khmer_det_repvit_db_latest.safetensors     # as published (Hub commit 5c9f76c)
+  sha256: 94d633392ec4be36b679a272bcb741524e003649bcdcbaf3b81ab77a504dca2c
+  onnx: khmer_det_repvit_db_latest.onnx   # the default (torch-free) install runs this
+  onnx_sha256: 5d5fdff595e644a4cc95f8c23e8b455a7a003293b35476212c0b09820c304440
+input:
+  limit_side_len: 1280
+  min_side_len: 736
+  mean: [0.485, 0.456, 0.406]
+  std: [0.229, 0.224, 0.225]
+architecture:
+  algorithm: DB_mobile
+  Backbone: {name: RepSVTR_det}
+  Neck: {name: RSEFPN, out_channels: 96, shortcut: true}
+  Head: {name: DBHead, k: 50}
+postprocess: {name: DBPostProcess, thresh: 0.3, box_thresh: 0.5, max_candidates: 1000,
+  unclip_ratio: 1.5, score_mode: slow, use_dilation: false}

anaksor/data/khmer_dict_v3.txt

@@ -0,0 +1,211 @@
+ក
+ខ
+គ
+ឃ
+ង
+ច
+ឆ
+ជ
+ឈ
+ញ
+ដ
+ឋ
+ឌ
+ឍ
+ណ
+ត
+ថ
+ទ
+ធ
+ន
+ប
+ផ
+ព
+ភ
+ម
+យ
+រ
+ល
+វ
+ឝ
+ឞ
+ស
+ហ
+ឡ
+អ
+ឥ
+ឦ
+ឧ
+ឨ
+ឩ
+ឪ
+ឫ
+ឬ
+ឭ
+ឮ
+ឯ
+ឰ
+ឱ
+ឲ
+ឳ
+ា
+ិ
+ី
+ឹ
+ឺ
+ុ
+ូ
+ួ
+ើ
+ឿ
+ៀ
+េ
+ែ
+ៃ
+ោ
+ៅ
+ំ
+ះ
+ៈ
+៉
+៊
+់
+៌
+៍
+៎
+៏
+័
+៑
+្
+៓
+។
+៕
+៖
+ៗ
+៙
+៚
+៛
+ៜ
+៝
+០
+១
+២
+៣
+៤
+៥
+៦
+៧
+៨
+៩
+!
+"
+#
+$
+%
+&
+'
+(
+)
+*
++
+,
+-
+.
+/
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+:
+;
+<
+=
+>
+?
+@
+A
+B
+C
+D
+E
+F
+G
+H
+I
+J
+K
+L
+M
+N
+O
+P
+Q
+R
+S
+T
+U
+V
+W
+X
+Y
+Z
+[
+\
+]
+^
+_
+`
+a
+b
+c
+d
+e
+f
+g
+h
+i
+j
+k
+l
+m
+n
+o
+p
+q
+r
+s
+t
+u
+v
+w
+x
+y
+z
+{
+|
+}
+~
+«
+»
+°
+×
+÷
+–
+—
+•
+€
+£
+¥
+©
+®
+±
+§
+№
+→
+✓

anaksor/data/rec_svtrv2_v3.yml

@@ -0,0 +1,51 @@
+# Khmer text-line recognition v3: SVTRv2 encoder, CTC head (+ SMTR head, unused by default).
+# Architecture block as in Alpha-OCR configs/svtrv2_khmer_v3.yml (training-only keys are ignored).
+# The SMTR vocabulary is </s> + the dict + <s>, <INF>, <INB>, <pad>; the CTC one is blank + the dict.
+name: rec-svtrv2-v3
+kind: rec
+weights:
+  file: khmer_rec_svtrv2_v3.safetensors
+  sha256: null            # set when the release weights are published
+  onnx: khmer_rec_svtrv2_v3.onnx      # tools/export_onnx.py; the default (torch-free) install runs this
+  onnx_sha256: null
+charset:
+  dict: khmer_dict_v3.txt
+  use_space_char: true
+  max_text_length: 140
+input:                    # same shape buckets as training (RatioDataSetTVResize, no padding)
+  base_h: 48
+  base_shape: [[64, 64], [96, 48], [144, 48], [192, 48]]
+  max_ratio: 20
+crop:                     # margin around each detected line, fraction of its height
+  pad_y: 0.22             # live sweep, 12 pages, CER: 0.04 -> 3.9%, 0.22/0.20 -> 0.34%
+  pad_x: 0.20
+architecture:
+  model_type: rec
+  algorithm: BGPD
+  in_channels: 3
+  Transform:
+  Encoder:
+    name: SVTRv2LNConvTwo33
+    use_pos_embed: false
+    out_channels: 256
+    dims: [96, 192, 384]
+    depths: [3, 6, 3]
+    num_heads: [3, 6, 12]
+    mixer: [[Conv, Conv, Conv], [Conv, Conv, Conv, FGlobal, Global, Global], [Global, Global, Global]]
+    local_k: [[5, 5], [5, 5], [-1, -1]]
+    sub_k: [[1, 1], [2, 1], [-1, -1]]
+    last_stage: false
+    feat2d: true
+  Decoder:
+    name: GTCDecoder
+    infer_gtc: true
+    detach: false
+    gtc_decoder:
+      name: SMTRDecoder
+      num_layer: 1
+      ds: true
+      max_len: 140
+      next_mode: true
+      sub_str_len: 5
+    ctc_decoder:
+      name: RCTCDecoder

anaksor/data/rec_svtrv2_v4.yml

@@ -0,0 +1,57 @@
+# Khmer text-line recognition v4: SVTRv2-B encoder, CTC head (+ SMTR head, unused by default;
+# the counting head only trained the encoder). Made with tools/spec_from_config.py from the
+# training config published as rec_latest.yml at Sokheng/anaksor-ocr, tag rec-v4-det-v3.
+name: rec-svtrv2-v4
+kind: rec
+weights:
+  file: khmer_rec_svtrv2_latest.safetensors     # as published (Hub commit 5c9f76c)
+  sha256: 2aead8175ed67ce8492363ae8664963643f1dc2108b74c663ed224b2b648f41d
+  onnx: khmer_rec_svtrv2_latest.onnx      # the default (torch-free) install runs this
+  onnx_sha256: afe130ebd442b35cff553efc91213219d33afc5f1c1883a0ddad402a8d1c96ec
+charset: {dict: khmer_dict_v3.txt, use_space_char: true, max_text_length: 140}
+input:
+  base_h: 48
+  base_shape:
+  - [64, 64]
+  - [96, 48]
+  - [144, 48]
+  - [192, 48]
+  max_ratio: 20
+crop: {pad_y: 0.16, pad_x: 0.1}
+architecture:
+  model_type: rec
+  algorithm: BGPD
+  in_channels: 3
+  Encoder:
+    name: SVTRv2LNConvTwo33
+    use_pos_embed: false
+    dims: [128, 256, 384]
+    depths: [6, 6, 6]
+    num_heads: [4, 8, 12]
+    mixer:
+    - [Conv, Conv, Conv, Conv, Conv, Conv]
+    - [Conv, Conv, FGlobal, Global, Global, Global]
+    - [Global, Global, Global, Global, Global, Global]
+    kernel_sizes:
+    - [5, 5, 5, 5, 5, 5]
+    - [5, 5, 5, 5, 5, 5]
+    - [-1]
+    num_convs:
+    - [1, 1, 1, 1, 1, 1]
+    - [1, 1, 1, 1, 1, 1]
+    - [-1]
+    pope_bias: true
+    sub_k:
+    - [1, 1]
+    - [2, 1]
+    - [-1, -1]
+    last_stage: false
+    feat2d: true
+  Decoder:
+    name: GTCCountDecoder
+    infer_gtc: true
+    detach: false
+    count_head: {max_len: 140, count_bins: 32, num_layer: 2}
+    gtc_decoder: {name: SMTRDecoder, num_layer: 1, ds: true, max_len: 140, next_mode: true,
+      sub_str_len: 5}
+    ctc_decoder: {name: RCTCDecoder}

anaksor/db_boxes.py

@@ -0,0 +1,105 @@
+"""Text boxes from a DBNet probability map.
+
+DBNet's inference (Liao et al., AAAI 2020, sec. 3.4): binarise the shrunk-text probability
+map with a fixed threshold, take each connected region's contour, discard regions whose
+mean probability is low, then *unclip*: grow the region by D = A * r / L (A area, L
+perimeter, r the unclip ratio) to undo the shrinking applied to the training labels.
+Each grown region becomes the minimum-area rotated rectangle around it, scaled back to
+the source image.
+"""
+
+from __future__ import annotations
+
+import cv2
+import numpy as np
+import pyclipper
+from shapely.geometry import Polygon
+
+MIN_SIDE = 3  # px on the probability map: smaller regions are noise
+
+
+def boxes_from_map(
+    prob: np.ndarray,
+    src_hw: tuple[int, int],
+    thresh: float = 0.3,
+    box_thresh: float = 0.5,
+    unclip_ratio: float = 1.5,
+    max_candidates: int = 1000,
+    score_mode: str = "slow",
+    dilate: bool = False,
+) -> tuple[list[np.ndarray], list[float]]:
+    """Quads (4 x 2 int32, clockwise from top-left) and scores of the text regions in
+    ``prob`` (H x W, float), in the coordinates of a ``src_hw`` image.
+
+    score_mode "slow" averages the probability inside the region's contour, "fast"
+    inside its rotated rectangle."""
+    if score_mode not in ("slow", "fast"):
+        raise ValueError(f"score_mode must be slow or fast, not {score_mode!r}")
+    mask = (prob > thresh).astype(np.uint8)
+    if dilate:
+        mask = cv2.dilate(mask, np.ones((2, 2), np.uint8))
+    height, width = mask.shape
+    src_h, src_w = np.float64(src_hw[0]), np.float64(src_hw[1])
+    contours, _ = cv2.findContours(mask * 255, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
+
+    boxes, scores = [], []
+    for contour in contours[:max_candidates]:
+        quad, short = _min_area_quad(contour)
+        if short < MIN_SIDE:
+            continue
+        score = _mean_in_polygon(prob, contour.reshape(-1, 2)) if score_mode == "slow" else _mean_in_box(prob, quad)
+        if score < box_thresh:
+            continue
+        grown = _unclip(quad, unclip_ratio)
+        if len(grown) != 1:  # vanished or split: not a text line
+            continue
+        quad, short = _min_area_quad(np.array(grown).reshape(-1, 1, 2))
+        if short < MIN_SIDE + 2:
+            continue
+        quad[:, 0] = np.clip(np.round(quad[:, 0] / width * src_w), 0, src_w)
+        quad[:, 1] = np.clip(np.round(quad[:, 1] / height * src_h), 0, src_h)
+        boxes.append(quad.astype(np.int32))
+        scores.append(score)
+    return boxes, scores
+
+
+def _min_area_quad(points) -> tuple[np.ndarray, float]:
+    """Corners of the minimum-area rectangle around ``points`` in the order top-left,
+    top-right, bottom-right, bottom-left, and its shorter side."""
+    rect = cv2.minAreaRect(points)
+    pts = sorted(cv2.boxPoints(rect), key=lambda p: p[0])  # two left corners, then two right
+    tl, bl = (pts[0], pts[1]) if pts[1][1] > pts[0][1] else (pts[1], pts[0])
+    tr, br = (pts[2], pts[3]) if pts[3][1] > pts[2][1] else (pts[3], pts[2])
+    return np.array([tl, tr, br, bl]), min(rect[1])
+
+
+def _unclip(quad: np.ndarray, ratio: float) -> list:
+    """Offset the polygon outwards by area * ratio / perimeter (round joins)."""
+    poly = Polygon(quad)
+    offset = pyclipper.PyclipperOffset()
+    offset.AddPath(quad, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
+    return offset.Execute(poly.area * ratio / poly.length)
+
+
+def _mean_in_polygon(prob: np.ndarray, pts: np.ndarray) -> float:
+    h, w = prob.shape
+    x0, x1 = np.clip(pts[:, 0].min(), 0, w - 1), np.clip(pts[:, 0].max(), 0, w - 1)
+    y0, y1 = np.clip(pts[:, 1].min(), 0, h - 1), np.clip(pts[:, 1].max(), 0, h - 1)
+    return _masked_mean(prob, pts, x0, x1, y0, y1)
+
+
+def _mean_in_box(prob: np.ndarray, quad: np.ndarray) -> float:
+    h, w = prob.shape
+    x0 = np.clip(np.floor(quad[:, 0].min()).astype(np.int32), 0, w - 1)
+    x1 = np.clip(np.ceil(quad[:, 0].max()).astype(np.int32), 0, w - 1)
+    y0 = np.clip(np.floor(quad[:, 1].min()).astype(np.int32), 0, h - 1)
+    y1 = np.clip(np.ceil(quad[:, 1].max()).astype(np.int32), 0, h - 1)
+    return _masked_mean(prob, quad, x0, x1, y0, y1)
+
+
+def _masked_mean(prob, pts, x0, x1, y0, y1) -> float:
+    """Mean of ``prob`` inside polygon ``pts``, over the window [y0..y1, x0..x1]."""
+    mask = np.zeros((y1 - y0 + 1, x1 - x0 + 1), np.uint8)
+    local = (pts - np.array([x0, y0], dtype=pts.dtype)).astype(np.int32)
+    cv2.fillPoly(mask, local.reshape(1, -1, 2), 1)
+    return cv2.mean(prob[y0 : y1 + 1, x0 : x1 + 1], mask)[0]