diff --git a/.claude/skills/datapipe-examples/SKILL.md b/.claude/skills/datapipe-examples/SKILL.md index 7e43b7c7..6170ec0d 100644 --- a/.claude/skills/datapipe-examples/SKILL.md +++ b/.claude/skills/datapipe-examples/SKILL.md @@ -19,13 +19,14 @@ Each `examples/*` pipeline has a dedicated setup skill — pick by what it does: | `e2e_template/image_keypoints` | YOLO-pose keypoints + Label Studio → train → FiftyOne | **setup-e2e-template** | | `sam_cvat` | SAM3 text-prompt boxes+masks → CVAT pre-annotations | **setup-sam-cvat** | | `detection_tags` | YOLO detection + **tags** (per-scenario metrics), no Label Studio / FiftyOne, GT injected | **setup-detection-tags** | +| `ocr` | Multi-LLM structured OCR (OpenAI/Gemini/Qwen) on passport/id-doc images → FiftyOne | **setup-ocr** | ## Ask first — don't assume (only the unresolved) Clarify what's not obvious before acting — don't spin up services or pull data you don't need. Ask only the unresolved: -- **Demo or your own data?** · **Provision Postgres/services or reuse existing?** (e2e ships `docker compose`; embedder/sam need external Postgres) +- **Demo or your own data?** · **Provision Postgres/services or reuse existing?** (e2e / detection_tags ship `docker compose`; embedder/sam/ocr need external Postgres) - **Which Postgres + which database** for `DB_URL`? Never point it at an existing DB or drop in a `localhost` default without confirming. -- **Reuse an existing venv / `uv` env, or create a fresh one?** · **Which GPU?** (SAM3 >8 GB; DINOv2/YOLO ~8 GB) · **Annotation backend up?** (LS for e2e / external CVAT for sam) -- **Surface stage logs or run quiet?** · **Per-tag scenario metrics** (retrain new case, old vs new)? → the `detection_tags` example / **setup-detection-tags** +- **Reuse an existing venv / `uv` env, or create a fresh one?** · **Which GPU?** (SAM3 >8 GB; DINOv2/YOLO ~8 GB; OCR needs none — LLM APIs) · **Annotation backend up?** (LS for e2e / external CVAT for sam) +- **OCR engines + API keys?** (`OCR_ENGINES` → openai/gemini/qwen) · **Surface stage logs or run quiet?** · **Per-tag scenario metrics** (retrain new case, old vs new)? → the `detection_tags` example / **setup-detection-tags** ## How to work Read the setup, then propose a short plan and get a go-ahead before touching anything. Prepare `.env` and **pause for the user to verify it** before running. Run each stage with its logs shown and, after each, say what you did and what changed — don't run the pipeline silently. Trust the status table (`*_training_status`/`brain_status`), not the exit code. If a stage fails and the cause isn't clear from the normal logs, re-run it with `datapipe --debug … run` (or `--debug-sql` for SQL errors) sent to a file and `grep`ped, rather than dumping the verbose output inline. @@ -36,12 +37,13 @@ prompt / label config, the label field, and any class filter. A mismatch runs fi useful results. The per-example skill lists the exact knobs. ## Universal prerequisites (every example) -1. **PostgreSQL** at `DB_URL`. embedder + sam_cvat need an EXTERNAL Postgres (they don't start one); - e2e_template bundles Postgres in its `docker compose`. Empty DB is fine; tables auto-create via +1. **PostgreSQL** at `DB_URL`. embedder + sam_cvat + ocr need an EXTERNAL Postgres (they don't start one); + e2e_template / detection_tags bundle Postgres in `docker compose`. Empty DB is fine; tables auto-create via `datapipe db create-all`. `.env.example` default `...postgres:postgres@localhost:5432/postgres`. External quick start: `docker run -d -e POSTGRES_PASSWORD=postgres -e POSTGRES_USER=postgres -e POSTGRES_DB=postgres -p 5432:5432 postgres:16` (k8s pod / no docker → conda postgres, `initdb` as non-root with `--locale=C`.) 2. **A GPU big enough** (`nvidia-smi`): DINOv2/YOLO fit ~8 GB; SAM3 wants >8 GB (OOM'd on an 8 GB Pascal in our tests). + OCR needs no GPU (calls LLM APIs; needs API keys for enabled engines). 3. **Annotation backend, if used:** e2e_template ships Label Studio in its `docker compose`; sam_cvat needs an external CVAT you provide (URL + creds + a project whose labels match its config). 4. **`uv` + Python ≥3.10,<3.13.** Each example has a `pyproject.toml` → `uv sync` (cu124 torch pinned, @@ -52,7 +54,8 @@ useful results. The per-example skill lists the exact knobs. ## HF auth (only for GATED models) Needed for `dinov3_*` (embedder, off by default) and `sam3` (sam_cvat). Accept the model license on its HF page, then `export HF_TOKEN=...` with **gated-repo read access**. A plain token → `403`, often -masked as `LocalEntryNotFoundError: check your connection`. Public models (dinov2, YOLO) need no token. +masked as `LocalEntryNotFoundError: check your connection`. Public models (dinov2, YOLO) and the OCR +demo HF dataset (`ud-synthetic/printed-usa-passports`) need no token. ## Universal run shape ```bash diff --git a/.claude/skills/setup-ocr/SKILL.md b/.claude/skills/setup-ocr/SKILL.md new file mode 100644 index 00000000..bc6e4026 --- /dev/null +++ b/.claude/skills/setup-ocr/SKILL.md @@ -0,0 +1,81 @@ +--- +name: setup-ocr +description: > + Use when working in examples/ocr, or running / debugging the OCR→FiftyOne + datapipe example (multi-LLM structured OCR on passport/id-doc images). +--- + +# ocr (LLM OCR → FiftyOne) + +This skill = run multi-engine LLM OCR on YOUR passport/id-doc images and publish structured JSON to FiftyOne. The HF dataset (`ud-synthetic/printed-usa-passports`) is just a smoke-test; the real goal is your own images — set the knobs below first. + +**Ask first — don't assume (only the unresolved):** demo (HF passports) or your own images? **which Postgres + which database** for `DB_URL` — never point it at an existing DB or drop in a `localhost` default without confirming; reuse an existing venv / `uv` env or create a fresh one? **which OCR engines** (`OCR_ENGINES` → subset of `openai,gemini,qwen`; each needs its API key)? **which model names** (defaults below — keep or change in `config.py` `ENGINE_REGISTRY`)? **cost cap** (`HF_LIMIT` in demo mode; local mode = all images × all engines)? **custom output schema** (default `IdDocument` in `config.py` or user edits `OUTPUT_MODEL`)? surface stage logs or run quiet? + +**How to work:** read the setup, then propose a short plan and get a go-ahead before touching anything. Prepare `.env` and **pause for the user to verify it** before running. Run each stage with its logs shown and, after each, say what you did and what changed — don't run the pipeline silently. If a stage fails and the cause isn't clear from the normal logs, re-run it with `datapipe --debug … run` (or `--debug-sql` for SQL errors); debug is very verbose, so send it to a file and `grep` it (e.g. `datapipe --debug run > /tmp/dp_debug.log 2>&1; grep -nEi "error|traceback" /tmp/dp_debug.log`) rather than dumping it inline. + +## Run on YOUR data +- **Your images:** `LOCAL_IMAGES_DIR=/path` (recursive `.jpg/.jpeg/.png/.webp/.heic`; file stem → `image_id`). + Local wins when the dir exists and has images (`use_local_images()` in `config.py`); otherwise HF fallback. +- **Or HF:** `HF_DATASET_NAME`, `HF_SPLIT`, `HF_LIMIT` (caps download **and** LLM calls in demo mode). +- **Engines:** `OCR_ENGINES=openai,gemini` (comma list). Model names + inference kwargs live in `config.py` + `ENGINE_REGISTRY` (not `.env`) — **ask which models to use**, then edit `ENGINE_REGISTRY[..]["model"]` if needed. + Current defaults (verify against `config.py`): `openai` → `gpt-5.4-nano`, `gemini` → `gemini-3.5-flash`, + `qwen` → `qwen3.7-plus`. +- **FiftyOne dataset:** `FIFTYONE_DATASET_NAME` (default `datapipe_ocr`). +- **Optional:** edit `OUTPUT_MODEL` in `config.py` — `ocr_prompt()` auto-follows the pydantic schema. + +## Prerequisites +Stages (labels match `app.py`): `ingest` (list images + engines) → `ocr` (LLM inference, images × engines +cross-join, `chunk_size=2`) → `fiftyone` (per-engine `{engine_id}_ocr` StringField via Caption Viewer). +- **External PostgreSQL** at `DB_URL` (SQLAlchemy URL). Not started by the example; an empty DB is fine — + `datapipe db create-all` auto-creates tables. `.env.example` default `...postgres@localhost:5432/postgres`. +- **LLM API keys** for each enabled engine: `OPENAI_API_KEY`, `GEMINI_API_KEY` (or `GOOGLE_API_KEY`), `QWEN_API_KEY`. + **Geoblock:** `openai` / `gemini` call the **native** OpenAI / Google APIs — blocked from Russian IPs. Workarounds + (proxy/router) are possible outside this example; from RU, prefer `qwen` (DashScope) or run from a non-blocked network. +- **FiftyOne must run on the SAME machine** — samples live in local FiftyOne/MongoDB; remote → empty panels. +- **Caption Viewer plugin** (required for readable JSON in the App): + `fiftyone plugins download https://github.com/harpreetsahota204/caption_viewer` +- **No GPU needed** — OCR calls LLM APIs, not local inference. **`uv` + Python 3.10–3.12** → `uv sync` + (pins CPU `torch==2.6.0` for transitive cv-pipeliner; OCR itself doesn't use it). +- Default HF dataset (`ud-synthetic/printed-usa-passports`) is public — no `HF_TOKEN` required. + +## Quick demo to verify setup +Skip if you have data. Leave `LOCAL_IMAGES_DIR` unset → HF fallback (`HF_DATASET_NAME`=ud-synthetic/printed-usa-passports, +`HF_LIMIT=10` in `.env.example` keeps LLM cost low); run §Run as-is to confirm install/DB/API keys/FiftyOne. + +## Run (from examples/ocr) +```bash +cp .env.example .env # DB_URL, API keys, OCR_ENGINES, LOCAL_IMAGES_DIR or HF_*, FIFTYONE_DATASET_NAME +uv sync && source .venv/bin/activate # else prefix each command with `uv run` +fiftyone plugins download https://github.com/harpreetsahota204/caption_viewer # once per venv +datapipe db create-all && datapipe run +# one stage: datapipe step --labels=stage=ingest run (ingest|ocr|fiftyone) +``` +Run from `examples/ocr/` (`load_dotenv()` in `app.py` — wrong cwd breaks `.env`). + +## Cost note +OCR runs a **cross-product** of images × enabled engines. Demo: `HF_LIMIT` caps both. Local: every image hits +every engine — confirm image count and `OCR_ENGINES` before a full run. + +## View in FiftyOne (same machine) +```bash +fiftyone app launch --remote --address 0.0.0.0 --port 5151 --wait -1 # venv active; SSH-forward 5151 +``` +Open `FIFTYONE_DATASET_NAME`. Per engine: open a sample → add panel → **Caption Viewer** → field +`openai_ocr` / `gemini_ocr` / `qwen_ocr`. Open multiple panels to compare engines side by side. + +## Success criteria +- `ocr_results` table has a row per `(image_id, engine_id)` pair with populated `output_json`. +- FiftyOne dataset has samples with `{engine_id}_ocr` fields containing structured JSON. +- Caption Viewer panel renders the JSON readably. + +## Troubleshooting (may already be fixed — verify against current files) +- **No images found** → empty/missing `LOCAL_IMAGES_DIR` and HF misconfigured; or local dir has no supported suffixes. +- **No OCR engines enabled** → `OCR_ENGINES` empty or unknown IDs (must match `ENGINE_REGISTRY` keys in `config.py`). +- **API key errors** → missing/wrong key for an enabled engine (`OPENAI_API_KEY`, `GEMINI_API_KEY`, `QWEN_API_KEY`). +- **OpenAI/Gemini timeout / region / connection refused from RU** → geoblock on native APIs (see Prerequisites); use `qwen` or a non-blocked network/proxy. +- **Runaway LLM cost** → forgot `HF_LIMIT` in demo mode, or large local folder with all three engines enabled. +- **FiftyOne empty / mongod bind errors** → old `~/.fiftyone` datadir (FCV mismatch / port bind). Use a fresh + dir: `FIFTYONE_DATABASE_DIR=/tmp/fo_db`. +- **JSON not pretty in App** → Caption Viewer plugin not installed (see Run). +- **Schema mismatch after editing `OUTPUT_MODEL`** → re-run `datapipe step --labels=stage=ocr run` so engines pick up the new prompt. diff --git a/examples/ocr/.env.example b/examples/ocr/.env.example new file mode 100644 index 00000000..17ebfc54 --- /dev/null +++ b/examples/ocr/.env.example @@ -0,0 +1,12 @@ +DB_URL=postgresql+psycopg2://postgres:postgres@localhost:5432/postgres +OCR_ENGINES=openai,gemini,qwen +OPENAI_API_KEY=replace-me +GEMINI_API_KEY=replace-me +QWEN_API_KEY=replace-me +# local mode: use when dir exists and contains images (optional HF fallback) +# LOCAL_IMAGES_DIR=./data/images +HF_DATASET_NAME=ud-synthetic/printed-usa-passports +HF_SPLIT=train +HF_LIMIT=10 +DATA_DIR=./data +FIFTYONE_DATASET_NAME=datapipe_ocr diff --git a/examples/ocr/.gitignore b/examples/ocr/.gitignore new file mode 100644 index 00000000..8fce6030 --- /dev/null +++ b/examples/ocr/.gitignore @@ -0,0 +1 @@ +data/ diff --git a/examples/ocr/README.md b/examples/ocr/README.md new file mode 100644 index 00000000..11416e3c --- /dev/null +++ b/examples/ocr/README.md @@ -0,0 +1,104 @@ +# OCR + FiftyOne Pipeline + +**Claude Code skill:** `/setup-ocr` — auto-loads when relevant, or invoke directly. + +Datapipe example for: +- ingesting passport / id-doc images (local folder or Hugging Face dataset) +- running OCR with multiple LLM engines (OpenAI, Gemini) +- publishing pydantic `OUTPUT_MODEL` JSON to FiftyOne StringFields +- viewing results with the Caption Viewer plugin + +## Data sources + +Supports **local folder** or **Hugging Face dataset** fallback. + +- Local: `LOCAL_IMAGES_DIR` with `.jpg`/`.jpeg`/`.png`/`.webp`/`.heic` inside (recursive scan) +- HF fallback when local dir unset, missing, or empty: `HF_DATASET_NAME` (default `ud-synthetic/printed-usa-passports`) +- HF images cached under `DATA_DIR//` + +Use `HF_LIMIT` to cap HF download size and LLM cost (ignored for local dirs). + +## Engine switch + +Single source of truth: `config.ENGINE_REGISTRY` + `.env` `OCR_ENGINES`. + +```text +OCR_ENGINES=openai,gemini,qwen +``` + +FiftyOne tables and pipeline steps are generated from the same enabled engine list, so config and tables stay aligned. + +Model names and inference kwargs (LLM `temperature`, etc.) live in `ENGINE_REGISTRY` inside `config.py`, not in `.env`. + +## Prerequisites + +### PostgreSQL + +Set `DB_URL` in `.env`: + +```text +postgresql+psycopg2://user:password@host:port/dbname +``` + +Create tables (use a fresh DB or drop old tables after schema changes): + +```shell +datapipe db create-all +``` + +### API keys + +- `OPENAI_API_KEY` for `openai` +- `GEMINI_API_KEY` for `gemini` +- `QWEN_API_KEY` for `qwen` (DashScope compatible OpenAI API) + +### FiftyOne App + +First - after creating virtual env install the community plugin for readable JSON in the FiftyOne modal: + +```shell +fiftyone plugins download https://github.com/harpreetsahota204/caption_viewer +``` + +Launch on the same machine as the pipeline: + +```shell +fiftyone app launch --remote --address 0.0.0.0 --port 5151 --wait -1 +``` + +To visualize JSON structured output in the App: open a sample → add panel → **Caption Viewer** → select `openai_ocr`, `gemini_ocr`, or `qwen_ocr`. + +## Run + +1. Copy env: `cp .env.example .env` +2. Edit `.env` (DB, API keys, dataset, limits) +3. Install deps: `uv sync` +4. From this folder: + +```shell +datapipe run +``` + +Or run stages by label: + +```shell +datapipe step --labels=stage=ingest run +datapipe step --labels=stage=ocr run +datapipe step --labels=stage=fiftyone run +``` + +## Visualize in FiftyOne + +Dataset name: `FIFTYONE_DATASET_NAME` (default `datapipe_ocr`). + +| Engine | FiftyOne field | How to view | +|--------|----------------|-------------| +| `openai` | `openai_ocr` | Caption Viewer panel (pretty JSON) | +| `gemini` | `gemini_ocr` | Caption Viewer panel (pretty JSON) | +| `qwen` | `qwen_ocr` | Caption Viewer panel (pretty JSON) | + +Open multiple Caption Viewer panels to compare engine fields side by side. + +## Customize output schema + +Edit `OUTPUT_MODEL` in `config.py` (default: passport fields `name`, `surname`, `birthdate`). All engines use it as output schema. OCR instruction text auto-generated via `ocr_prompt()`. diff --git a/examples/ocr/app.py b/examples/ocr/app.py new file mode 100644 index 00000000..d27f3c8f --- /dev/null +++ b/examples/ocr/app.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +import logging + +from dotenv import load_dotenv + +load_dotenv() +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") + +from datapipe.compute import Catalog, DatapipeApp, Pipeline +from datapipe.datatable import DataStore +from datapipe.step.batch_generate import BatchGenerate +from datapipe.step.batch_transform import BatchTransform + +import data +import steps +from config import DBCONN, ENABLED_ENGINES + +pipeline_steps = [ + BatchGenerate( + steps.list_images, + outputs=[data.images_tbl], + labels=[("stage", "ingest")], + ), + BatchGenerate( + steps.list_engines, + outputs=[data.engines_tbl], + labels=[("stage", "ingest")], + ), + BatchTransform( + func=steps.ocr_inference, + inputs=[data.images_tbl, data.engines_tbl], + outputs=[data.ocr_results_tbl], + transform_keys=["image_id", "engine_id"], + chunk_size=2, + labels=[("stage", "ocr")], + ), +] + +for engine_id in ENABLED_ENGINES: + pipeline_steps.append( + BatchTransform( + func=steps.make_build_image_data(engine_id), + inputs=[data.images_tbl, data.ocr_results_tbl], + outputs=[data.fo_tables[engine_id]], + transform_keys=["image_id"], + labels=[("stage", "fiftyone")], + ) + ) + +pipeline = Pipeline(pipeline_steps) + +ds = DataStore(DBCONN, create_meta_table=True) +app = DatapipeApp(ds, Catalog({}), pipeline) diff --git a/examples/ocr/config.py b/examples/ocr/config.py new file mode 100644 index 00000000..4d47d194 --- /dev/null +++ b/examples/ocr/config.py @@ -0,0 +1,103 @@ +from __future__ import annotations + +import os +from datetime import date +from pathlib import Path +from typing import Any + +from datapipe.store.database import DBConn +from pydantic import BaseModel, Field + +DATA_DIR = Path(os.environ.get("DATA_DIR", "./data")).resolve() + +_local_images_dir_raw = os.environ.get("LOCAL_IMAGES_DIR") +LOCAL_IMAGES_DIR = Path(_local_images_dir_raw).resolve() if _local_images_dir_raw else None + +IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp", ".heic"} + +HF_DATASET_NAME = os.environ.get("HF_DATASET_NAME", "ud-synthetic/printed-usa-passports") +HF_SPLIT = os.environ.get("HF_SPLIT", "train") +HF_LIMIT = int(os.environ.get("HF_LIMIT", "50")) + +OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "") +GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", os.environ.get("GOOGLE_API_KEY", "")) +QWEN_API_KEY = os.environ.get("QWEN_API_KEY", "") + +DB_URL = os.environ.get("DB_URL") +DBCONN = DBConn(DB_URL, "public") + +FIFTYONE_DATASET_NAME = os.environ.get("FIFTYONE_DATASET_NAME", "datapipe_ocr") + + +class IdDocument(BaseModel): + """Default OCR output for passport / id-doc images.""" + + name: str = Field(description="Given name(s) as printed on the document") + surname: str = Field(description="Family name / surname as printed on the document") + birthdate: date | None = Field(default=None, description="Date of birth in ISO format YYYY-MM-DD") + + +OUTPUT_MODEL = IdDocument + + +def ocr_prompt(model: type[BaseModel] = OUTPUT_MODEL) -> str: + """Build OCR instruction text from the pydantic output schema.""" + header = (model.__doc__ or "Extract fields from the document image.").strip() + field_lines = [ + f"- {field_name}: {field_info.description or field_name}" + for field_name, field_info in model.model_fields.items() + ] + return f"{header}\n\nReturn JSON matching these fields:\n" + "\n".join(field_lines) + + +def _parse_ocr_engines() -> list[str]: + raw = os.environ.get("OCR_ENGINES", "openai,gemini") + return [engine_id.strip() for engine_id in raw.split(",") if engine_id.strip()] + + +ENGINE_REGISTRY: dict[str, dict[str, Any]] = { + "openai": { + "provider": "openai", + "model": "gpt-5.4-nano", + "init_kwargs": { + "temperature": 0, + }, + }, + "gemini": { + "provider": "gemini", + "model": "gemini-3.5-flash", + "init_kwargs": { + "temperature": 0, + }, + }, + "qwen": { + "provider": "qwen", + "model": "qwen3.7-plus", + "client_kwargs": { + "base_url": "https://dashscope-intl.aliyuncs.com/compatible-mode/v1", + }, + "init_kwargs": { + "extra_body": {"enable_thinking": False}, + "response_format": {"type": "json_object"}, + "temperature": 0, + }, + }, +} + +ENABLED_ENGINES = [engine_id for engine_id in _parse_ocr_engines() if engine_id in ENGINE_REGISTRY] + + +def fo_text_field(engine_id: str) -> str: + return f"{engine_id}_ocr" + + +def _local_dir_has_images(local_dir: Path | None) -> bool: + if local_dir is None or not local_dir.exists(): + return False + return any( + path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES for path in local_dir.rglob("*") + ) + + +def use_local_images() -> bool: + return _local_dir_has_images(LOCAL_IMAGES_DIR) diff --git a/examples/ocr/data.py b/examples/ocr/data.py new file mode 100644 index 00000000..5a36dd55 --- /dev/null +++ b/examples/ocr/data.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +from cv_pipeliner.utils.fiftyone import FifyOneSession +from datapipe.compute import Table +from datapipe.store.database import TableStoreDB +from datapipe_ml.utils.image_data_stores import FiftyOneImagesDataTableStore +from sqlalchemy import Column, String, Text + +from config import DBCONN, ENABLED_ENGINES, FIFTYONE_DATASET_NAME, fo_text_field + +fo_session = FifyOneSession() + +images_tbl = Table( + name="images", + store=TableStoreDB( + dbconn=DBCONN, + name="images", + data_sql_schema=[ + Column("image_id", String(255), primary_key=True), + Column("image_path", String(2048)), + ], + create_table=True, + ), +) + +engines_tbl = Table( + name="engines", + store=TableStoreDB( + dbconn=DBCONN, + name="engines", + data_sql_schema=[ + Column("engine_id", String(64), primary_key=True), + Column("provider", String(64)), + Column("model", String(255)), + ], + create_table=True, + ), +) + +ocr_results_tbl = Table( + name="ocr_results", + store=TableStoreDB( + dbconn=DBCONN, + name="ocr_results", + data_sql_schema=[ + Column("image_id", String(255), primary_key=True), + Column("engine_id", String(64), primary_key=True), + Column("output_json", Text), + Column("full_text", Text), + ], + create_table=True, + ), +) + +fo_tables: dict[str, Table] = {} +for engine_id in ENABLED_ENGINES: + fo_tables[engine_id] = Table( + name=f"fo_{engine_id}", + store=FiftyOneImagesDataTableStore( + dataset=FIFTYONE_DATASET_NAME, + fo_session=fo_session, + additional_info_keys_in_sample=[fo_text_field(engine_id)], + rm_only_fo_fields=True, + primary_schema=[Column("image_id", String(255), primary_key=True)], + ), + ) diff --git a/examples/ocr/engines.py b/examples/ocr/engines.py new file mode 100644 index 00000000..7379045d --- /dev/null +++ b/examples/ocr/engines.py @@ -0,0 +1,176 @@ +from __future__ import annotations + +import base64 +import logging +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from config import ( + ENGINE_REGISTRY, + GEMINI_API_KEY, + OPENAI_API_KEY, + OUTPUT_MODEL, + QWEN_API_KEY, + ocr_prompt, +) + +logger = logging.getLogger(__name__) + + +@dataclass +class OcrResult: + output_json: str + full_text: str + + +def _encode_image_base64(image_path: str) -> str: + return base64.b64encode(Path(image_path).read_bytes()).decode("utf-8") + + +def _guess_mime_type(image_path: str) -> str: + suffix = Path(image_path).suffix.lower() + if suffix in {".jpg", ".jpeg"}: + return "image/jpeg" + if suffix == ".png": + return "image/png" + if suffix == ".webp": + return "image/webp" + if suffix == ".heic": + return "image/heic" + return "image/jpeg" + + +def _engine_config(engine_id: str) -> dict[str, Any]: + return ENGINE_REGISTRY[engine_id] + + +def _document_from_output(document: Any) -> Any: + if isinstance(document, OUTPUT_MODEL): + return document + return OUTPUT_MODEL.model_validate(document) + + +def _ocr_result_from_document(document: Any) -> OcrResult: + document = _document_from_output(document) + output_json = document.model_dump_json(indent=2) + return OcrResult(output_json=output_json, full_text=output_json) + + +def run_openai(image_path: str, engine_id: str = "openai", **_kwargs: Any) -> OcrResult: + if not OPENAI_API_KEY: + raise ValueError("OPENAI_API_KEY is required for openai engine") + + from openai import OpenAI + + engine_cfg = _engine_config(engine_id) + init_kwargs = engine_cfg.get("init_kwargs", {}) + prompt = ocr_prompt(OUTPUT_MODEL) + + client = OpenAI(api_key=OPENAI_API_KEY) + mime = _guess_mime_type(image_path) + b64 = _encode_image_base64(image_path) + response = client.responses.parse( + model=engine_cfg["model"], + input=[ + { + "role": "user", + "content": [ + {"type": "input_text", "text": prompt}, + { + "type": "input_image", + "image_url": f"data:{mime};base64,{b64}", + }, + ], + } + ], + text_format=OUTPUT_MODEL, + **init_kwargs, + ) + document = response.output_parsed + if document is None: + raise ValueError("OpenAI structured parse returned empty output") + return _ocr_result_from_document(document) + + +def run_gemini(image_path: str, engine_id: str = "gemini", **_kwargs: Any) -> OcrResult: + if not GEMINI_API_KEY: + raise ValueError("GEMINI_API_KEY (or GOOGLE_API_KEY) is required for gemini engine") + + from google import genai + from google.genai import types + + engine_cfg = _engine_config(engine_id) + init_kwargs = engine_cfg.get("init_kwargs", {}) + prompt = ocr_prompt(OUTPUT_MODEL) + + client = genai.Client(api_key=GEMINI_API_KEY) + image_bytes = Path(image_path).read_bytes() + mime = _guess_mime_type(image_path) + response = client.models.generate_content( + model=engine_cfg["model"], + contents=[ + types.Part.from_bytes(data=image_bytes, mime_type=mime), + prompt, + ], + config=types.GenerateContentConfig( + response_mime_type="application/json", + response_schema=OUTPUT_MODEL, + **init_kwargs, + ), + ) + document = getattr(response, "parsed", None) + if document is None: + document = OUTPUT_MODEL.model_validate_json(response.text or "{}") + return _ocr_result_from_document(document) + + +def run_qwen(image_path: str, engine_id: str = "qwen", **_kwargs: Any) -> OcrResult: + if not QWEN_API_KEY: + raise ValueError("QWEN_API_KEY is required for qwen engine") + + from openai import OpenAI + + engine_cfg = _engine_config(engine_id) + client_kwargs = engine_cfg.get("client_kwargs", {}) + init_kwargs = engine_cfg.get("init_kwargs", {}) + prompt = ocr_prompt(OUTPUT_MODEL) + + client = OpenAI(api_key=QWEN_API_KEY, **client_kwargs) + mime = _guess_mime_type(image_path) + b64 = _encode_image_base64(image_path) + response = client.chat.completions.create( + model=engine_cfg["model"], + messages=[ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": f"data:{mime};base64,{b64}"}, + }, + {"type": "text", "text": prompt}, + ], + } + ], + **init_kwargs, + ) + content = response.choices[0].message.content + if not content: + raise ValueError("Qwen completion returned empty content") + document = OUTPUT_MODEL.model_validate_json(content) + return _ocr_result_from_document(document) + + +ENGINE_RUNNERS = { + "openai": run_openai, + "gemini": run_gemini, + "qwen": run_qwen, +} + + +def run_engine(engine_id: str, image_path: str) -> OcrResult: + runner = ENGINE_RUNNERS.get(engine_id) + if runner is None: + raise KeyError(f"Unknown OCR engine: {engine_id}") + return runner(image_path=image_path, engine_id=engine_id) diff --git a/examples/ocr/pyproject.toml b/examples/ocr/pyproject.toml new file mode 100644 index 00000000..13238efc --- /dev/null +++ b/examples/ocr/pyproject.toml @@ -0,0 +1,26 @@ +[project] +name = "ocr-example" +version = "0" +requires-python = ">=3.10,<3.13" +dependencies = [ + "datapipe-ml[fiftyone]", + "python-dotenv>=1.0", + "stringzilla==4.4.0", # needed for albumentations + "fiftyone==1.17.0", + "openai>=1.0", + "google-genai>=1.0", + "datasets>=3.0", + "pydantic>=2.0", + "torch==2.6.0", # transitive via cv-pipeliner; pin CPU wheel (not used by OCR) +] + +[[tool.uv.index]] +name = "pytorch-cpu" +url = "https://download.pytorch.org/whl/cpu" +explicit = true + +[tool.uv.sources] +datapipe-ml = { path = "../../libs/datapipe-ml", editable = true } +datapipe-core = { path = "../../libs/datapipe-core", editable = true } +cv-pipeliner = { git = "https://github.com/epoch8/cv-pipeliner", rev = "5724f8d54e4df64013fad85d41129799bc143293" } +torch = { index = "pytorch-cpu" } diff --git a/examples/ocr/steps.py b/examples/ocr/steps.py new file mode 100644 index 00000000..dea89ddc --- /dev/null +++ b/examples/ocr/steps.py @@ -0,0 +1,183 @@ +from __future__ import annotations + +import json +import logging +from collections.abc import Callable, Iterator +from pathlib import Path + +import pandas as pd +from cv_pipeliner import ImageData +from PIL import Image + +from config import ( + DATA_DIR, + ENGINE_REGISTRY, + ENABLED_ENGINES, + HF_DATASET_NAME, + HF_LIMIT, + HF_SPLIT, + IMAGE_SUFFIXES, + LOCAL_IMAGES_DIR, + fo_text_field, + use_local_images, +) +from engines import run_engine + +logger = logging.getLogger(__name__) + + +def _dataset_cache_dir(dataset_name: str) -> Path: + return DATA_DIR / dataset_name.replace("/", "__") + + +def _save_pil_image(image: Image.Image, path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + image.convert("RGB").save(path, format="JPEG", quality=95) + + +def _load_hf_split(dataset_name: str, split: str): + from datasets import load_dataset + + return load_dataset(dataset_name, split=split) + + +def _list_images_from_local_dir(local_dir: Path) -> list[dict[str, str]]: + files = sorted( + path + for path in local_dir.rglob("*") + if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES + ) + return [ + { + "image_id": path.stem, + "image_path": str(path.resolve()), + } + for path in files + ] + + +def _hf_sample_stem(item: dict, idx: int) -> str: + for key in ("sample_id", "Sample ID", "passport_id", "Passport ID", "id", "filename"): + if key in item and item[key] is not None: + return str(item[key]).replace("/", "_").replace(" ", "_") + return f"item_{idx:06d}" + + +def _materialize_images() -> list[dict[str, str]]: + if use_local_images(): + assert LOCAL_IMAGES_DIR is not None + records = _list_images_from_local_dir(LOCAL_IMAGES_DIR) + logger.info("Using local images from %s: %d images", LOCAL_IMAGES_DIR, len(records)) + return records + + ds = _load_hf_split(HF_DATASET_NAME, HF_SPLIT) + save_dir = _dataset_cache_dir(HF_DATASET_NAME) + records: list[dict[str, str]] = [] + + for idx, item in enumerate(ds): + if idx >= HF_LIMIT: + break + + filename = item.get("filename") + if filename: + stem = Path(str(filename)).stem + out_path = save_dir / str(filename) + else: + stem = _hf_sample_stem(item, idx) + out_path = save_dir / f"{stem}.jpg" + + if not out_path.exists(): + image = item.get("image") + if image is None: + logger.warning("Skipping sample %s: no image field", idx) + continue + _save_pil_image(image, out_path) + + records.append( + { + "image_id": stem, + "image_path": str(out_path.resolve()), + } + ) + + logger.info("Using HF dataset %s: %d images", HF_DATASET_NAME, len(records)) + return records + + +def list_images() -> Iterator[pd.DataFrame]: + records = _materialize_images() + if not records: + raise ValueError( + "No images found. Set LOCAL_IMAGES_DIR with images or configure HF_DATASET_NAME in .env" + ) + yield pd.DataFrame(records, columns=["image_id", "image_path"]) + + +def list_engines() -> Iterator[pd.DataFrame]: + rows = [ + { + "engine_id": engine_id, + "provider": ENGINE_REGISTRY[engine_id]["provider"], + "model": ENGINE_REGISTRY[engine_id]["model"], + } + for engine_id in ENABLED_ENGINES + ] + if not rows: + raise ValueError("No OCR engines enabled. Set OCR_ENGINES in .env") + yield pd.DataFrame(rows) + + +def ocr_inference(images_df: pd.DataFrame, engines_df: pd.DataFrame) -> pd.DataFrame: + if images_df.empty or engines_df.empty: + return pd.DataFrame(columns=["image_id", "engine_id", "output_json", "full_text"]) + + merged = images_df.merge(engines_df, how="cross") + if merged.empty: + return pd.DataFrame(columns=["image_id", "engine_id", "output_json", "full_text"]) + + records = [] + for _, row in merged.iterrows(): + result = run_engine(engine_id=row["engine_id"], image_path=row["image_path"]) + records.append( + { + "image_id": row["image_id"], + "engine_id": row["engine_id"], + "output_json": result.output_json, + "full_text": result.full_text, + } + ) + + return pd.DataFrame(records) + + +def _build_image_data(image_path: str, engine_id: str, ocr_row: pd.Series) -> ImageData: + text_field = fo_text_field(engine_id) + output_json = ocr_row.get("output_json") or ocr_row.get("full_text") or "" + if isinstance(output_json, dict): + output_json = json.dumps(output_json, indent=2) + return ImageData( + image_path=image_path, + bboxes_data=[], + additional_info={text_field: str(output_json)}, + ) + + +def make_build_image_data(engine_id: str) -> Callable[[pd.DataFrame, pd.DataFrame], pd.DataFrame]: + def build_image_data(images_df: pd.DataFrame, ocr_results_df: pd.DataFrame) -> pd.DataFrame: + if images_df.empty or ocr_results_df.empty: + return pd.DataFrame(columns=["image_id", "image_data"]) + + engine_results = ocr_results_df[ocr_results_df["engine_id"] == engine_id] + if engine_results.empty: + return pd.DataFrame(columns=["image_id", "image_data"]) + + merged = images_df.merge(engine_results, on="image_id", how="inner") + records = [] + for _, row in merged.iterrows(): + image_data = _build_image_data(str(row["image_path"]), engine_id, row) + records.append({"image_id": row["image_id"], "image_data": image_data}) + + return pd.DataFrame(records, columns=["image_id", "image_data"]) + + build_image_data.__name__ = f"build_image_data_{engine_id}" + return build_image_data