uv-scripts/ocr

1# /// script2# requires-python = ">=3.10"3# dependencies = [4#     "paddlepaddle-gpu>=3.0.0",5#     "paddleocr>=3.7.0",6#     "paddlex[ocr]>=3.7.0",7#     "opencv-contrib-python-headless",8#     "datasets>=3.1.0",9#     "huggingface-hub",10#     "pillow",11#     "numpy",12#     "tqdm",13# ]14#15# [tool.uv]16# # PaddleOCR/PaddleX pull in opencv-contrib-python (full) which needs system17# # libGL.so.1 — not present in the slim uv-on-bookworm image used by HF Jobs.18# # Swap to the headless cv2 variant (same `import cv2`, no GUI deps). A matching19# # importlib.metadata patch in main() makes paddlex recognise the headless name.20# override-dependencies = [21#     "opencv-contrib-python ; python_version < '0'",22#     "opencv-python ; python_version < '0'",23# ]24#25# [[tool.uv.index]]26# name = "paddle"27# url = "https://www.paddlepaddle.org.cn/packages/stable/cu126/"28# explicit = true29#30# [tool.uv.sources]31# paddlepaddle-gpu = { index = "paddle" }32#33# [tool.hf-jobs]34# flavor = "t4-small"35# secrets = ["HF_TOKEN"]36# ///37"""38OCR images with PP-OCRv6 — a lightweight detection+recognition pipeline from39PaddlePaddle. Three tiers from **1.5M to 34.5M parameters**.40 41Unlike the VLM-based OCR recipes here, PP-OCRv6 is a **classical det+rec pipeline**42that outputs **plain text** (not markdown). At 1.5M-34.5M params it's far smaller43than the VLM OCRs and runs on a cheap t4-small GPU.44 45Model tiers (pick with `--model-tier`):46  tiny    1.5M params  (0.4M det + 1.1M rec)  49 languages, ~73% recognition47  small   7.7M params  (2.5M det + 5.3M rec)  50 languages, ~81% recognition48  medium  34.5M params (22M det + 19M rec)     50 languages, ~83% recognition49 50All tiers are Apache 2.0 licensed. Runs via PaddleOCR's default Paddle engine51(`paddle_static`) — same proven header pattern as `pp-doclayout.py`.52 53HF Jobs examples (flavor and secrets come from the [tool.hf-jobs] header,54which needs `hf` CLI 1.32+):55 56    # Tiny on a cheap GPU57    hf jobs uv run \\58        https://huggingface.co/datasets/uv-scripts/ocr/raw/main/pp-ocrv6.py \\59        INPUT_DATASET OUTPUT_DATASET \\60        --model-tier tiny --max-samples 561 62    # Medium on a small GPU (recommended for quality)63    hf jobs uv run \\64        https://huggingface.co/datasets/uv-scripts/ocr/raw/main/pp-ocrv6.py \\65        INPUT_DATASET OUTPUT_DATASET \\66        --model-tier medium --max-samples 1067 68Models: PaddlePaddle/PP-OCRv6_<tier>_det + PP-OCRv6_<tier>_rec69Blog: https://huggingface.co/blog/PaddlePaddle/pp-ocrv670"""71 72import argparse73import io74import json75import logging76import os77import sys78import time79from dataclasses import dataclass80from datetime import datetime, timezone81from pathlib import Path82from typing import Any, Dict, Iterator, List, Optional, Tuple, Union83 84import numpy as np85from PIL import Image, UnidentifiedImageError86from tqdm.auto import tqdm87 88logging.basicConfig(level=logging.INFO)89logger = logging.getLogger(__name__)90 91 92# ---------------------------------------------------------------------------93# Constants94# ---------------------------------------------------------------------------95 96TIER_MODELS = {97    "tiny":   ("PP-OCRv6_tiny_det",   "PP-OCRv6_tiny_rec"),98    "small":  ("PP-OCRv6_small_det",  "PP-OCRv6_small_rec"),99    "medium": ("PP-OCRv6_medium_det", "PP-OCRv6_medium_rec"),100}101 102TIER_PARAMS = {103    "tiny":   "1.5M (0.4M det + 1.1M rec)",104    "small":  "7.7M (2.5M det + 5.3M rec)",105    "medium": "34.5M (22M det + 19M rec)",106}107 108TIER_LANGUAGES = {109    "tiny":   "49 languages (zh, zh-Hant, en + 46 Latin-script — no Japanese)",110    "small":  "50 languages (zh, zh-Hant, en, ja + 46 Latin-script)",111    "medium": "50 languages (zh, zh-Hant, en, ja + 46 Latin-script)",112}113 114TIER_REC = {115    "tiny":   73.5,116    "small":  81.3,117    "medium": 83.2,118}119 120BUCKET_PREFIX = "hf://buckets/"121 122IMAGE_EXTENSIONS = {123    ".jpg", ".jpeg", ".png", ".tif", ".tiff", ".webp", ".bmp", ".jp2", ".j2k",124}125 126 127# ---------------------------------------------------------------------------128# URL helpers129# ---------------------------------------------------------------------------130 131def is_bucket_url(s: str) -> bool:132    return s.startswith(BUCKET_PREFIX)133 134 135def parse_bucket_url(url: str) -> Tuple[str, str]:136    if not is_bucket_url(url):137        raise ValueError(f"Not a bucket URL: {url}")138    rest = url[len(BUCKET_PREFIX):].strip("/")139    parts = rest.split("/", 2)140    if len(parts) < 2:141        raise ValueError(f"Bucket URL must include namespace and bucket name: {url}")142    bucket_id = f"{parts[0]}/{parts[1]}"143    prefix = parts[2] if len(parts) > 2 else ""144    return bucket_id, prefix145 146 147# ---------------------------------------------------------------------------148# Image helpers149# ---------------------------------------------------------------------------150 151def to_pil(image: Union[Image.Image, Dict[str, Any], str, bytes]) -> Image.Image:152    if isinstance(image, Image.Image):153        return image.convert("RGB")154    if isinstance(image, dict) and "bytes" in image:155        return Image.open(io.BytesIO(image["bytes"])).convert("RGB")156    if isinstance(image, (bytes, bytearray)):157        return Image.open(io.BytesIO(image)).convert("RGB")158    if isinstance(image, str):159        return Image.open(image).convert("RGB")160    raise ValueError(f"Unsupported image type: {type(image)}")161 162 163def pil_to_array(pil_img: Image.Image) -> np.ndarray:164    return np.asarray(pil_img, dtype=np.uint8)165 166 167# ---------------------------------------------------------------------------168# Result extraction169# ---------------------------------------------------------------------------170 171def extract_text(result: Any) -> Tuple[str, List[Dict[str, Any]]]:172    """Pull text and per-line details from a PaddleOCR predict result.173 174    Returns (concatenated_text, per_line_details) where per_line_details is175    a list of dicts with keys: text, score, bbox (4-point detection polygon as176    [[x1,y1],[x2,y2],[x3,y3],[x4,y4]] in input-image pixel coordinates).177    """178    payload = result.json if hasattr(result, "json") else result179    res = payload.get("res", payload) if isinstance(payload, dict) else {}180    rec_texts = res.get("rec_texts", []) or []181    rec_scores = res.get("rec_scores", []) or []182    dt_polys = res.get("dt_polys", []) or []183 184    # Concatenate reading-order text lines (PaddleOCR returns them in order)185    text = "\n".join(rec_texts)186 187    per_line = []188    for i, t in enumerate(rec_texts):189        entry = {"text": t}190        if i < len(rec_scores):191            entry["score"] = float(rec_scores[i])192        if i < len(dt_polys):193            entry["bbox"] = [[float(c) for c in point] for point in dt_polys[i]]194        per_line.append(entry)195 196    return text, per_line197 198 199# ---------------------------------------------------------------------------200# Sources201# ---------------------------------------------------------------------------202 203@dataclass204class SourceItem:205    key: str206    image: Optional[Image.Image]207    extras: Dict[str, Any]208 209 210def iter_dataset_images(211    dataset_id: str,212    image_column: str,213    split: str,214    shuffle: bool,215    seed: int,216    max_samples: Optional[int],217):218    from datasets import load_dataset219 220    logger.info(f"Loading dataset: {dataset_id} (split={split})")221    ds = load_dataset(dataset_id, split=split)222 223    if image_column not in ds.column_names:224        raise ValueError(225            f"Column '{image_column}' not found. Available: {ds.column_names}"226        )227 228    if shuffle:229        logger.info(f"Shuffling with seed {seed}")230        ds = ds.shuffle(seed=seed)231    if max_samples:232        ds = ds.select(range(min(max_samples, len(ds))))233        logger.info(f"Limited to {len(ds)} samples")234 235    total = len(ds)236 237    def gen() -> Iterator[SourceItem]:238        failed = 0239        for i in range(total):240            try:241                row = ds[i]242                image = to_pil(row[image_column])243            except (UnidentifiedImageError, OSError) as e:244                # Still yield a placeholder so the output row stays aligned with245                # the source row (the dataset sink writes results positionally).246                failed += 1247                logger.warning(248                    f"Unreadable image at row {i}: {type(e).__name__}: {e} "249                    f"— writing empty result"250                )251                yield SourceItem(key=f"row-{i:08d}", image=None, extras={"failed": True})252                continue253            yield SourceItem(key=f"row-{i:08d}", image=image, extras={})254        if failed:255            logger.info(f"{failed} unreadable image(s) written as empty results")256 257    return gen(), total, ds258 259 260SOURCE_PATHS_SNAPSHOT = "_source_paths.json"261 262 263def _bucket_snapshot_path(output_url: str) -> Tuple[str, str]:264    out_bucket_id, out_prefix = parse_bucket_url(output_url)265    snapshot_key = (266        f"{out_prefix}/{SOURCE_PATHS_SNAPSHOT}".lstrip("/")267        if out_prefix268        else SOURCE_PATHS_SNAPSHOT269    )270    return out_bucket_id, snapshot_key271 272 273def iter_bucket_images(274    bucket_url: str,275    shuffle: bool,276    seed: int,277    max_samples: Optional[int],278    hf_token: Optional[str],279    output_url: Optional[str] = None,280) -> Tuple[Iterator[SourceItem], int]:281    from huggingface_hub import HfApi, HfFileSystem282 283    bucket_id, prefix = parse_bucket_url(bucket_url)284    fs = HfFileSystem(token=hf_token)285    base = f"{BUCKET_PREFIX}{bucket_id}/{prefix}".rstrip("/")286 287    snapshot_bucket_id: Optional[str] = None288    snapshot_key: Optional[str] = None289    cached_paths: Optional[List[str]] = None290 291    if output_url and is_bucket_url(output_url):292        snapshot_bucket_id, snapshot_key = _bucket_snapshot_path(output_url)293        snapshot_url = f"{BUCKET_PREFIX}{snapshot_bucket_id}/{snapshot_key}"294        try:295            with fs.open(snapshot_url, "rb") as f:296                snapshot = json.load(f)297            mismatches = []298            if snapshot.get("source_url") != bucket_url:299                mismatches.append(300                    f"source_url ({snapshot.get('source_url')!r} vs {bucket_url!r})"301                )302            if snapshot.get("shuffle") != shuffle:303                mismatches.append(f"shuffle ({snapshot.get('shuffle')} vs {shuffle})")304            if shuffle and snapshot.get("seed") != seed:305                mismatches.append(f"seed ({snapshot.get('seed')} vs {seed})")306            if snapshot.get("max_samples") != max_samples:307                mismatches.append(308                    f"max_samples ({snapshot.get('max_samples')} vs {max_samples})"309                )310            if mismatches:311                logger.warning(312                    "Existing snapshot params differ from this run ("313                    + "; ".join(mismatches)314                    + "); ignoring snapshot and re-listing."315                )316            else:317                cached_paths = snapshot["paths"]318                logger.info(319                    f"Reusing existing snapshot of {len(cached_paths)} source paths "320                    f"(written {snapshot.get('created_at', 'unknown')})"321                )322        except FileNotFoundError:323            pass324        except Exception as e:325            logger.warning(f"Could not read existing snapshot ({e}); re-listing.")326 327    if cached_paths is not None:328        all_paths = cached_paths329    else:330        logger.info(f"Listing images under {base}")331        all_paths = []332        try:333            for entry in fs.find(base, detail=False):334                ext = Path(entry).suffix.lower()335                if ext in IMAGE_EXTENSIONS:336                    all_paths.append(entry)337        except FileNotFoundError as e:338            raise ValueError(f"Bucket prefix not found: {base}") from e339 340        if not all_paths:341            raise ValueError(342                f"No image files (any of {sorted(IMAGE_EXTENSIONS)}) under {base}"343            )344 345        all_paths.sort()346        if shuffle:347            rng = np.random.default_rng(seed)348            rng.shuffle(all_paths)349        if max_samples:350            all_paths = all_paths[:max_samples]351 352        if snapshot_bucket_id is not None and snapshot_key is not None:353            api = HfApi(token=hf_token)354            payload = {355                "source_url": bucket_url,356                "shuffle": shuffle,357                "seed": seed,358                "max_samples": max_samples,359                "created_at": datetime.now(timezone.utc).isoformat(),360                "paths": all_paths,361            }362            api.batch_bucket_files(363                snapshot_bucket_id,364                add=[(json.dumps(payload).encode(), snapshot_key)],365                token=hf_token,366            )367            logger.info(368                f"Wrote source-path snapshot ({len(all_paths)} paths) to "369                f"hf://buckets/{snapshot_bucket_id}/{snapshot_key}"370            )371 372    total = len(all_paths)373    logger.info(f"Found {total} images in bucket")374 375    def key_for(path: str) -> str:376        return path377 378    def gen() -> Iterator[SourceItem]:379        skipped = 0380        for path in all_paths:381            try:382                with fs.open(path, "rb") as f:383                    data = f.read()384                image = to_pil(data)385            except (UnidentifiedImageError, OSError) as e:386                skipped += 1387                logger.warning(388                    f"Skipping unreadable image {path}: {type(e).__name__}: {e}"389                )390                continue391            yield SourceItem(key=key_for(path), image=image, extras={})392        if skipped:393            logger.info(f"Skipped {skipped} unreadable image(s) total")394 395    return gen(), total396 397 398# ---------------------------------------------------------------------------399# Sinks400# ---------------------------------------------------------------------------401 402class DatasetRepoSink:403    def __init__(404        self,405        repo_id: str,406        *,407        hf_token: Optional[str],408        private: bool,409        config: Optional[str],410        create_pr: bool,411        source_id: str,412        original_dataset=None,413        output_column: str = "markdown",414        overwrite: bool = False,415    ):416        self.repo_id = repo_id417        self.hf_token = hf_token418        self.private = private419        self.config = config420        self.create_pr = create_pr421        self.source_id = source_id422        self.original_dataset = original_dataset423        self.output_column = output_column424        self.overwrite = overwrite425        self._texts: List[str] = []426        self._blocks: List[str] = []427 428    @property429    def kind(self) -> str:430        return "dataset"431 432    def already_done(self) -> set:433        return set()434 435    def write(self, key: str, text: str, blocks: List[Dict[str, Any]]) -> None:436        self._texts.append(text)437        self._blocks.append(json.dumps(blocks, ensure_ascii=False))438 439    def finalize(self, tier: str, det_model: str, rec_model: str, args_dict: Dict[str, Any]) -> None:440        from datasets import Dataset441 442        if self.original_dataset is not None:443            if len(self._texts) != len(self.original_dataset):444                logger.warning(445                    f"Text count ({len(self._texts)}) != dataset rows "446                    f"({len(self.original_dataset)}); padding with empty strings."447                )448                while len(self._texts) < len(self.original_dataset):449                    self._texts.append("")450                    self._blocks.append("[]")451            # Guard again at save time in case the input column set changed under us.452            base = self.original_dataset453            clash = [c for c in (self.output_column, "pp_ocr_blocks") if c in base.column_names]454            if clash:455                if not self.overwrite:456                    raise ValueError(457                        f"Output column(s) {clash} already exist in the input dataset; "458                        f"pass a different --output-column, or --overwrite to replace them."459                    )460                logger.warning(f"--overwrite: replacing existing column(s) {clash}")461                base = base.remove_columns(clash)462            ds = base.add_column(self.output_column, self._texts)463            ds = ds.add_column("pp_ocr_blocks", self._blocks)464        else:465            if not self._texts:466                logger.warning("No rows produced; nothing to push.")467                return468            ds = Dataset.from_list([469                {"source_path": None, self.output_column: t, "pp_ocr_blocks": b}470                for t, b in zip(self._texts, self._blocks)471            ])472 473        inference_entry = build_inference_entry(tier, det_model, rec_model, args_dict)474 475        if "inference_info" in ds.column_names:476            logger.info("Updating existing inference_info column")477 478            def _update(example):479                try:480                    existing = (481                        json.loads(example["inference_info"])482                        if example["inference_info"]483                        else []484                    )485                except (json.JSONDecodeError, TypeError):486                    existing = []487                existing.append(inference_entry)488                return {"inference_info": json.dumps(existing)}489 490            ds = ds.map(_update)491        else:492            ds = ds.add_column(493                "inference_info", [json.dumps([inference_entry])] * len(ds)494            )495 496        logger.info(f"Pushing {len(ds)} rows to {self.repo_id}")497        push_kwargs = {498            "private": self.private,499            "token": self.hf_token,500            "max_shard_size": "500MB",501            "create_pr": self.create_pr,502            "commit_message": f"Add PP-OCRv6-{tier} OCR results ({len(ds)} samples)"503            + (f" [{self.config}]" if self.config else ""),504        }505        if self.config:506            push_kwargs["config_name"] = self.config507 508        max_retries = 3509        for attempt in range(1, max_retries + 1):510            try:511                if attempt > 1:512                    logger.warning("Disabling XET (fallback to HTTP upload)")513                    os.environ["HF_HUB_DISABLE_XET"] = "1"514                ds.push_to_hub(self.repo_id, **push_kwargs)515                break516            except Exception as e:517                logger.error(f"Upload attempt {attempt}/{max_retries} failed: {e}")518                if attempt == max_retries:519                    logger.error("All upload attempts failed.")520                    raise521                time.sleep(30 * (2 ** (attempt - 1)))522 523        from huggingface_hub import DatasetCard524 525        card = DatasetCard(526            create_dataset_card(527                source=self.source_id,528                tier=tier,529                det_model=det_model,530                rec_model=rec_model,531                num_samples=len(ds),532                processing_time=args_dict["processing_time"],533                engine=args_dict.get("engine", "paddle_static"),534                output_id=self.repo_id,535                output_column=self.output_column,536            )537        )538        card.push_to_hub(self.repo_id, token=self.hf_token)539        logger.info(f"Done: https://huggingface.co/datasets/{self.repo_id}")540 541 542class BucketShardSink:543    METADATA_FILE = "_metadata.json"544    SHARD_PATTERN = "shard-{:05d}.parquet"545 546    def __init__(547        self,548        bucket_url: str,549        *,550        hf_token: Optional[str],551        shard_size: int,552        resume: bool,553        source_id: str,554    ):555        from huggingface_hub import HfApi, HfFileSystem, create_bucket556 557        self.bucket_url = bucket_url558        self.bucket_id, self.prefix = parse_bucket_url(bucket_url)559        self.hf_token = hf_token560        self.shard_size = shard_size561        self.resume = resume562        self.source_id = source_id563 564        self._api = HfApi(token=hf_token)565        self._fs = HfFileSystem(token=hf_token)566 567        try:568            create_bucket(self.bucket_id, exist_ok=True, token=hf_token)569        except Exception as e:570            logger.warning(f"create_bucket('{self.bucket_id}') warning: {e}")571 572        self._buffer: List[Dict[str, Any]] = []573        self._next_shard_idx = self._discover_next_shard_idx()574        self._completed_keys = self._discover_completed_keys() if resume else set()575        if self._completed_keys:576            logger.info(577                f"Resume: found {len(self._completed_keys)} already-processed keys, will skip them"578            )579 580    @property581    def kind(self) -> str:582        return "bucket"583 584    def already_done(self) -> set:585        return self._completed_keys586 587    def _shard_path(self, idx: int) -> str:588        return self._join(self.SHARD_PATTERN.format(idx))589 590    def _join(self, name: str) -> str:591        return f"{self.prefix}/{name}".lstrip("/") if self.prefix else name592 593    def _list_existing_shards(self) -> List[str]:594        try:595            tree = self._api.list_bucket_tree(596                self.bucket_id, prefix=self.prefix or None, recursive=True597            )598        except Exception:599            return []600        shards: List[str] = []601        for item in tree:602            path = getattr(item, "path", None)603            ftype = getattr(item, "type", None)604            if not path or ftype not in (None, "file"):605                continue606            base = Path(path).name607            if base.startswith("shard-") and base.endswith(".parquet"):608                shards.append(path)609        return sorted(shards)610 611    def _discover_next_shard_idx(self) -> int:612        shards = self._list_existing_shards()613        max_idx = -1614        for s in shards:615            stem = Path(s).stem616            try:617                max_idx = max(max_idx, int(stem.split("-")[-1]))618            except ValueError:619                continue620        return max_idx + 1621 622    def _discover_completed_keys(self) -> set:623        import pyarrow.parquet as pq624 625        keys: set = set()626        for shard_path in self._list_existing_shards():627            full = f"{BUCKET_PREFIX}{self.bucket_id}/{shard_path}"628            try:629                with self._fs.open(full, "rb") as f:630                    table = pq.read_table(f, columns=["__source_key"])631                keys.update(table.column("__source_key").to_pylist())632            except Exception as e:633                logger.warning(f"Could not read keys from {shard_path}: {e}")634        return keys635 636    def _flush(self) -> None:637        if not self._buffer:638            return639        import pyarrow as pa640        import pyarrow.parquet as pq641 642        columns = ["__source_key", "text", "pp_ocr_blocks"]643        table_dict = {c: [row.get(c) for row in self._buffer] for c in columns}644        table = pa.Table.from_pydict(table_dict)645 646        buf = io.BytesIO()647        pq.write_table(table, buf, compression="zstd")648        data = buf.getvalue()649 650        shard_remote = self._shard_path(self._next_shard_idx)651        logger.info(652            f"Writing shard {self._next_shard_idx} ({len(self._buffer)} rows, "653            f"{len(data) / 1024 / 1024:.1f} MiB) to {shard_remote}"654        )655        self._api.batch_bucket_files(656            self.bucket_id, add=[(data, shard_remote)], token=self.hf_token657        )658        self._next_shard_idx += 1659        self._buffer.clear()660 661    def write(self, key: str, text: str, blocks: List[Dict[str, Any]]) -> None:662        row: Dict[str, Any] = {663            "__source_key": key,664            "text": text,665            "pp_ocr_blocks": json.dumps(blocks, ensure_ascii=False),666        }667        self._buffer.append(row)668        if len(self._buffer) >= self.shard_size:669            self._flush()670 671    def finalize(self, tier: str, det_model: str, rec_model: str, args_dict: Dict[str, Any]) -> None:672        self._flush()673        meta = {674            "model": f"PP-OCRv6_{tier}",675            "det_model": det_model,676            "rec_model": rec_model,677            "tier": tier,678            "engine": "paddle_static",679            "source": self.source_id,680            "shard_size": args_dict["shard_size"],681            "last_run_at": datetime.now(timezone.utc).isoformat(),682            "processing_time": args_dict.get("processing_time"),683        }684        meta_bytes = json.dumps(meta, indent=2).encode("utf-8")685        meta_path = self._join(self.METADATA_FILE)686        self._api.batch_bucket_files(687            self.bucket_id, add=[(meta_bytes, meta_path)], token=self.hf_token688        )689        logger.info(690            f"Done: https://huggingface.co/buckets/{self.bucket_id}"691            + (f"/{self.prefix}" if self.prefix else "")692        )693 694 695# ---------------------------------------------------------------------------696# inference_info + dataset card697# ---------------------------------------------------------------------------698 699def build_inference_entry(tier: str, det_model: str, rec_model: str, args_dict: Dict[str, Any]) -> Dict[str, Any]:700    return {701        "model_id": f"PaddlePaddle/PP-OCRv6_{tier}",702        "det_model": det_model,703        "rec_model": rec_model,704        "tier": tier,705        "params": TIER_PARAMS.get(tier, "unknown"),706        "rec_accuracy_pct": TIER_REC.get(tier),707        "languages": TIER_LANGUAGES.get(tier, ""),708        "engine": "paddle_static",709        # column_name is the key ocr-bench's column discovery reads; keep710        # output_column too for backward compat with existing outputs.711        "column_name": args_dict.get("output_column", "markdown"),712        "output_column": args_dict.get("output_column", "markdown"),713        "blocks_column": "pp_ocr_blocks",714        "timestamp": datetime.now(timezone.utc).isoformat(),715    }716 717 718def create_dataset_card(719    source: str,720    tier: str,721    det_model: str,722    rec_model: str,723    num_samples: int,724    processing_time: str,725    engine: str,726    output_id: str,727    output_column: str = "markdown",728) -> str:729    tier_display = tier.upper() if tier == "tiny" else tier.capitalize()730    if is_bucket_url(source):731        source_link = f"[{source}]({source})"732    else:733        source_link = f"[{source}](https://huggingface.co/datasets/{source})"734 735    return f"""---736tags:737- ocr738- text-recognition739- paddleocr740- pp-ocrv6741- uv-script742- generated743---744 745# OCR with PP-OCRv6 {tier_display}746 747Plain-text OCR results for images from {source_link}, produced by748PaddlePaddle's [PP-OCRv6](https://huggingface.co/collections/PaddlePaddle/pp-ocrv6)749{tier} pipeline ({TIER_PARAMS.get(tier, "unknown")}).750 751## Processing details752 753- **Source**: {source_link}754- **Model**: PP-OCRv6_{tier} ({det_model} + {rec_model})755- **Tier**: {tier} ({TIER_PARAMS.get(tier, "unknown")})756- **Recognition accuracy**: {TIER_REC.get(tier, "?"):.1f}%757- **Languages**: {TIER_LANGUAGES.get(tier, "")}758- **Engine**: {engine}759- **Samples**: {num_samples:,}760- **Processing time**: {processing_time}761- **Processing date**: {datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC")}762- **License**: Apache 2.0 (models)763 764## Schema765 766Each row contains the original columns plus:767 768- `{output_column}`: Plain text extracted from the image (reading-order concatenation of769  detected text lines, newline-separated).770- `pp_ocr_blocks`: JSON list, one dict per detected text line:771  ```json772  [773    {{774      "text": "recognized text",775      "score": 0.987,776      "bbox": [[x1, y1], [x2, y2], [x3, y3], [x4, y4]]777    }}778  ]779  ```780  `score` is the recognition confidence and `bbox` is the detection polygon781  (4-point quadrilateral in input-image pixel coordinates).782- `inference_info`: JSON list tracking every model applied to this dataset.783 784> **Note:** PP-OCRv6 is a classical detection+recognition pipeline, not a VLM.785> It outputs **plain text** rather than markdown. Per-line bounding boxes and786> confidence scores are available in `pp_ocr_blocks`.787 788## Usage789 790```python791import json792from datasets import load_dataset793 794ds = load_dataset("{output_id}", split="train")795print(ds[0]["{output_column}"])796for block in json.loads(ds[0]["pp_ocr_blocks"]):797    print(block["text"], block["score"])798```799 800## Reproduction801 802```bash803hf jobs uv run \\804    https://huggingface.co/datasets/uv-scripts/ocr/raw/main/pp-ocrv6.py \\805    {source} <output> --model-tier {tier}806```807 808Generated with [UV Scripts](https://huggingface.co/uv-scripts).809"""810 811 812# ---------------------------------------------------------------------------813# Main814# ---------------------------------------------------------------------------815 816def main(args: argparse.Namespace) -> None:817    from huggingface_hub import login818 819    start_time = datetime.now()820    hf_token = args.hf_token or os.environ.get("HF_TOKEN")821    if hf_token:822        login(token=hf_token)823 824    # ---------- tier → model names ----------825    if args.model_tier not in TIER_MODELS:826        raise ValueError(827            f"Invalid tier {args.model_tier!r}. Choose from: {list(TIER_MODELS)}"828        )829    det_model, rec_model = TIER_MODELS[args.model_tier]830    tier = args.model_tier831    logger.info(f"PP-OCRv6 {tier}: {det_model} + {rec_model}")832 833    # ---------- source ----------834    original_dataset = None835    if is_bucket_url(args.input_source):836        src_iter, total = iter_bucket_images(837            args.input_source,838            shuffle=args.shuffle,839            seed=args.seed,840            max_samples=args.max_samples,841            hf_token=hf_token,842            output_url=args.output_target,843        )844    else:845        src_iter, total, original_dataset = iter_dataset_images(846            args.input_source,847            image_column=args.image_column,848            split=args.split,849            shuffle=args.shuffle,850            seed=args.seed,851            max_samples=args.max_samples,852        )853        # Fail fast, before minutes of inference, if the output column would collide854        # with an existing input column (e.g. a 'text' ground-truth column). Writing855        # into it would either crash on push or silently overwrite the input data.856        # --overwrite opts in to replacing the existing column(s) instead of erroring.857        if original_dataset is not None:858            clash = [859                col860                for col in (args.output_column, "pp_ocr_blocks")861                if col in original_dataset.column_names862            ]863            if clash and not args.overwrite:864                logger.error(865                    f"Output column(s) {clash} already exist in the input dataset "866                    f"(columns: {original_dataset.column_names})."867                )868                logger.error(869                    "Choose a different --output-column, or pass --overwrite to replace them."870                )871                sys.exit(1)872            if clash:873                logger.warning(f"--overwrite: will replace existing column(s) {clash}")874 875    # ---------- sink ----------876    if is_bucket_url(args.output_target):877        sink: Union[BucketShardSink, DatasetRepoSink] = BucketShardSink(878            args.output_target,879            hf_token=hf_token,880            shard_size=args.shard_size,881            resume=not args.no_resume,882            source_id=args.input_source,883        )884    else:885        sink = DatasetRepoSink(886            args.output_target,887            hf_token=hf_token,888            private=args.private,889            config=args.config,890            create_pr=args.create_pr,891            source_id=args.input_source,892            original_dataset=original_dataset,893            output_column=args.output_column,894            overwrite=args.overwrite,895        )896 897    completed = sink.already_done()898 899    # ---------- model ----------900    # PaddleX gates `import cv2` at module load time on901    # `is_dep_available("opencv-contrib-python")`, which checks902    # `importlib.metadata.version(...)`. We ship `opencv-contrib-python-headless`903    # (same `cv2`, no system libGL.so.1 needed) — but that's a different904    # distribution name, so the gate fails and the OCR pipeline's `ocr` extra905    # check returns False. Patch the metadata lookup to alias the GUI cv2 distros906    # to the headless variant before importing paddleocr; this lets paddlex's own907    # `import cv2` succeed and `is_extra_available('ocr')` return True.908    import importlib.metadata as _metadata909 910    _orig_metadata_version = _metadata.version911 912    def _patched_metadata_version(dep_name):913        if dep_name in ("opencv-contrib-python", "opencv-python"):914            for headless_alias in (915                "opencv-contrib-python-headless",916                "opencv-python-headless",917            ):918                try:919                    return _orig_metadata_version(headless_alias)920                except _metadata.PackageNotFoundError:921                    continue922        return _orig_metadata_version(dep_name)923 924    _metadata.version = _patched_metadata_version925 926    # Silence the connectivity check for speed (not needed in a Job)927    os.environ.setdefault("PADDLE_PDX_DISABLE_MODEL_SOURCE_CHECK", "True")928 929    from paddleocr import PaddleOCR930 931    ocr = PaddleOCR(932        text_detection_model_name=det_model,933        text_recognition_model_name=rec_model,934        use_doc_orientation_classify=False,935        use_doc_unwarping=False,936        use_textline_orientation=False,937    )938 939    # ---------- loop ----------940    processed = 0941    skipped = 0942    errors = 0943    pbar = tqdm(src_iter, total=total, desc=f"PP-OCRv6 {tier}")944    for item in pbar:945        if item.key in completed:946            skipped += 1947            continue948        if item.extras.get("failed") or item.image is None:949            # Unreadable source image — write an empty result in position so the950            # output stays row-aligned with the source dataset.951            sink.write(item.key, "", [])952            errors += 1953            processed += 1954            continue955        try:956            arr = pil_to_array(item.image)957            result = ocr.predict(arr)958            if result:959                text, blocks = extract_text(result[0])960            else:961                text, blocks = "", []962        except Exception as e:963            logger.error(f"Error on {item.key}: {e}")964            text, blocks = "", []965            errors += 1966 967        sink.write(item.key, text, blocks)968        processed += 1969 970    duration = datetime.now() - start_time971    processing_time_str = f"{duration.total_seconds() / 60:.2f} min"972    logger.info(973        f"Processed {processed} (skipped {skipped}, errors {errors}) in {processing_time_str}"974    )975 976    args_dict = {977        "tier": tier,978        "det_model": det_model,979        "rec_model": rec_model,980        "engine": "paddle_static",981        "shard_size": args.shard_size,982        "processing_time": processing_time_str,983        "output_column": args.output_column,984    }985    sink.finalize(986        tier=tier,987        det_model=det_model,988        rec_model=rec_model,989        args_dict=args_dict,990    )991 992    if args.verbose:993        import importlib.metadata994 995        logger.info("--- Resolved package versions ---")996        for pkg in [997            "paddleocr",998            "paddlex",999            "paddlepaddle-gpu",1000            "huggingface-hub",1001            "datasets",1002            "pillow",1003            "numpy",1004        ]:1005            try:1006                logger.info(f"  {pkg}=={importlib.metadata.version(pkg)}")1007            except importlib.metadata.PackageNotFoundError:1008                logger.info(f"  {pkg}: not installed")1009        logger.info("--- End versions ---")1010 1011 1012# ---------------------------------------------------------------------------1013# CLI1014# ---------------------------------------------------------------------------1015 1016def build_parser() -> argparse.ArgumentParser:1017    p = argparse.ArgumentParser(1018        description="PP-OCRv6 OCR over an HF dataset or bucket of images.",1019        formatter_class=argparse.RawDescriptionHelpFormatter,1020    )1021    p.add_argument(1022        "input_source",1023        help="HF dataset id (namespace/dataset) OR hf://buckets/ns/bucket[/prefix]",1024    )1025    p.add_argument(1026        "output_target",1027        help="HF dataset id (namespace/dataset) OR hf://buckets/ns/bucket/run-name",1028    )1029    p.add_argument(1030        "--model-tier",1031        default="medium",1032        choices=list(TIER_MODELS),1033        help="PP-OCRv6 model tier: tiny (1.5M), small (7.7M), medium (34.5M). Default: medium.",1034    )1035    # Dataset-source-specific1036    p.add_argument(1037        "--image-column",1038        default="image",1039        help="Column containing images (dataset-repo source only, default: image)",1040    )1041    p.add_argument(1042        "--split",1043        default="train",1044        help="Dataset split (dataset-repo source only, default: train)",1045    )1046    p.add_argument(1047        "--max-samples", type=int, help="Limit number of samples (for testing)"1048    )1049    p.add_argument(1050        "--shuffle", action="store_true", help="Shuffle source before processing"1051    )1052    p.add_argument(1053        "--seed", type=int, default=42, help="Random seed for shuffle (default: 42)"1054    )1055    # Dataset-sink-specific1056    p.add_argument(1057        "--private", action="store_true", help="Private dataset output (dataset sink only)"1058    )1059    p.add_argument(1060        "--config",1061        help="Config/subset name when pushing to Hub (dataset sink only)",1062    )1063    p.add_argument(1064        "--create-pr",1065        action="store_true",1066        help="Create PR instead of direct push (dataset sink only)",1067    )1068    p.add_argument(1069        "--output-column",1070        default="markdown",1071        help=(1072            "Column name for the recognized text (dataset sink only, default: markdown). "1073            "Must not collide with an existing input column — many corpora already ship a "1074            "'text' ground-truth column, so 'text' would fail on push. Blocks always go to "1075            "'pp_ocr_blocks'."1076        ),1077    )1078    p.add_argument(1079        "--overwrite",1080        action="store_true",1081        help="Replace the output column(s) if they already exist in the input dataset "1082        "(default: error out to avoid clobbering an existing column).",1083    )1084    # Bucket-sink-specific1085    p.add_argument(1086        "--shard-size",1087        type=int,1088        default=256,1089        help="Rows per parquet shard for bucket sink (default: 256)",1090    )1091    p.add_argument(1092        "--no-resume",1093        action="store_true",1094        help="Disable resume scan when writing to a bucket sink",1095    )1096    # Auth + diagnostics1097    p.add_argument("--hf-token", help="Hugging Face API token (else uses HF_TOKEN env)")1098    p.add_argument(1099        "--verbose",1100        action="store_true",1101        help="Log resolved package versions at the end",1102    )1103    return p1104 1105 1106if __name__ == "__main__":1107    main(build_parser().parse_args())1108