From ff8864b0c7075ce64c8c6f288ab62a98d4117fbb Mon Sep 17 00:00:00 2001 From: korotas Date: Thu, 2 Jul 2026 16:22:33 +0200 Subject: [PATCH 1/6] initial commit - ocr seems working --- examples/ocr/.env.example | 17 ++ examples/ocr/.gitignore | 1 + examples/ocr/README.md | 116 ++++++++++++ examples/ocr/app.py | 64 +++++++ examples/ocr/config.py | 133 ++++++++++++++ examples/ocr/data.py | 101 +++++++++++ examples/ocr/engines.py | 195 ++++++++++++++++++++ examples/ocr/pyproject.toml | 30 ++++ examples/ocr/steps.py | 343 ++++++++++++++++++++++++++++++++++++ examples/ocr/viz.py | 78 ++++++++ 10 files changed, 1078 insertions(+) create mode 100644 examples/ocr/.env.example create mode 100644 examples/ocr/.gitignore create mode 100644 examples/ocr/README.md create mode 100644 examples/ocr/app.py create mode 100644 examples/ocr/config.py create mode 100644 examples/ocr/data.py create mode 100644 examples/ocr/engines.py create mode 100644 examples/ocr/pyproject.toml create mode 100644 examples/ocr/steps.py create mode 100644 examples/ocr/viz.py diff --git a/examples/ocr/.env.example b/examples/ocr/.env.example new file mode 100644 index 00000000..e7ce50cf --- /dev/null +++ b/examples/ocr/.env.example @@ -0,0 +1,17 @@ +DB_URL=postgresql+psycopg2://postgres:postgres@localhost:5432/postgres +OCR_ENGINES=paddle,openai,gemini +OPENAI_API_KEY=replace-me +GEMINI_API_KEY=replace-me +# local mode: use when dir exists and contains images (optional HF fallback per branch) +# LOCAL_UNSTRUCTURED_DIR=./data/unstructured +# LOCAL_STRUCTURED_DIR=./data/structured +# unstructured branch (paddle): free-form document scans +HF_DATASET_UNSTRUCTURED=MLap/Book-Scan-OCR +HF_SPLIT_UNSTRUCTURED=train +# structured branch (openai/gemini): id documents -> IdDocument(name,surname,birthdate) +HF_DATASET_STRUCTURED=ud-synthetic/printed-usa-passports +HF_SPLIT_STRUCTURED=train +HF_LIMIT_STRUCTURED=10 +HF_LIMIT_UNSTRUCTURED=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..5f0124ce --- /dev/null +++ b/examples/ocr/README.md @@ -0,0 +1,116 @@ +# OCR + FiftyOne Pipeline + +Datapipe example for: +- ingesting two Hugging Face datasets (unstructured book scans + structured passport images) +- running OCR with a config-driven engine switch (PaddleOCR v6, OpenAI, Gemini) +- publishing results to FiftyOne (native boxes for unstructured, JSON StringFields for structured) +- optional side-by-side composite renders for unstructured OCR + +## Data sources + +Each branch supports **local folder** or **Hugging Face dataset** fallback. + +### Unstructured (`paddle`) + +- Local: `LOCAL_UNSTRUCTURED_DIR` with `.jpg`/`.jpeg`/`.png`/`.webp`/`.heic` inside (recursive scan) +- HF fallback when local dir unset, missing, or empty: `HF_DATASET_UNSTRUCTURED` (default `MLap/Book-Scan-OCR`) +- HF images cached under `DATA_DIR//` +- Routed to engines with `engine_type=unstructured` + +### Structured (`openai`, `gemini`) + +- Local: `LOCAL_STRUCTURED_DIR` with supported images inside (recursive scan) +- HF fallback when local dir unset, missing, or empty: `HF_DATASET_STRUCTURED` (default `ud-synthetic/printed-usa-passports`) +- HF images cached under `DATA_DIR//` +- Routed to engines with `engine_type=structured` +- Both LLMs share the same pydantic schema `IdDocument(name, surname, birthdate)` from `config.py` + +Use `HF_LIMIT_UNSTRUCTURED` / `HF_LIMIT_STRUCTURED` 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=paddle,openai,gemini +``` + +FiftyOne tables and pipeline steps are generated from the same enabled engine list, so config and tables stay aligned. + +Model names and inference kwargs (PaddleOCR init, 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: + +```shell +datapipe db create-all +``` + +### API keys (structured engines) + +- `OPENAI_API_KEY` for `openai` +- `GEMINI_API_KEY` (or `GOOGLE_API_KEY`) for `gemini` + +### FiftyOne App + +Launch on the same machine as the pipeline: + +```shell +fiftyone app launch --remote --address 0.0.0.0 --port 5151 --wait -1 +``` + +### Caption Viewer plugin (structured JSON) + +Install the community plugin for readable JSON in the FiftyOne modal: + +```shell +fiftyone plugins download https://github.com/harpreetsahota204/caption_viewer +``` + +In the App: open a sample → add panel → **Caption Viewer** → select `openai_ocr` or `gemini_ocr`. + +## Run + +1. Copy env: `cp .env.example .env` +2. Edit `.env` (DB, API keys, datasets, 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 +datapipe step --labels=stage=composite run +``` + +## Visualize in FiftyOne + +Dataset name: `FIFTYONE_DATASET_NAME` (default `datapipe_ocr`). + +| Engine type | FiftyOne fields | How to view | +|-------------|-----------------|-------------| +| unstructured (`paddle`) | `paddle_boxes` detections, `paddle_ocr` text | Toggle detections overlay; box label = recognized text | +| structured (`openai`, `gemini`) | `openai_ocr`, `gemini_ocr` StringFields | Caption Viewer panel (pretty JSON) | + +Open multiple Caption Viewer panels to compare `openai_ocr` vs `gemini_ocr` side by side. + +Composite PNGs (unstructured only) are written to `data/composites/{engine_id}/` when running `stage=composite`. + +## Customize structured schema + +Edit `IdDocument` in `config.py`. Both OpenAI and Gemini use it as structured output schema. The OCR instruction text is generated automatically from that model via `structured_ocr_prompt()`. diff --git a/examples/ocr/app.py b/examples/ocr/app.py new file mode 100644 index 00000000..cfbee28f --- /dev/null +++ b/examples/ocr/app.py @@ -0,0 +1,64 @@ +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=8, + 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_steps.append( + BatchTransform( + func=steps.render_composite, + inputs=[data.images_tbl, data.ocr_results_tbl], + outputs=[data.composites_tbl], + transform_keys=["engine_id", "image_id"], + labels=[("stage", "composite")], + ) +) + +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..b99b3ea2 --- /dev/null +++ b/examples/ocr/config.py @@ -0,0 +1,133 @@ +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() +COMPOSITES_DIR = DATA_DIR / "composites" + +_local_unstructured_dir_raw = os.environ.get("LOCAL_UNSTRUCTURED_DIR") +LOCAL_UNSTRUCTURED_DIR = ( + Path(_local_unstructured_dir_raw).resolve() if _local_unstructured_dir_raw else None +) +_local_structured_dir_raw = os.environ.get("LOCAL_STRUCTURED_DIR") +LOCAL_STRUCTURED_DIR = Path(_local_structured_dir_raw).resolve() if _local_structured_dir_raw else None + +IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp", ".heic"} + +HF_DATASET_UNSTRUCTURED = os.environ.get("HF_DATASET_UNSTRUCTURED", "MLap/Book-Scan-OCR") +HF_SPLIT_UNSTRUCTURED = os.environ.get("HF_SPLIT_UNSTRUCTURED", "train") +HF_LIMIT_UNSTRUCTURED = int(os.environ.get("HF_LIMIT_UNSTRUCTURED", "50")) + +HF_DATASET_STRUCTURED = os.environ.get("HF_DATASET_STRUCTURED", "ud-synthetic/printed-usa-passports") +HF_SPLIT_STRUCTURED = os.environ.get("HF_SPLIT_STRUCTURED", "train") +HF_LIMIT_STRUCTURED = int(os.environ.get("HF_LIMIT_STRUCTURED", "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", "")) + +DB_URL = os.environ.get("DB_URL") +DBCONN = DBConn(DB_URL, "public") + +FIFTYONE_DATASET_NAME = os.environ.get("FIFTYONE_DATASET_NAME", "datapipe_ocr") + +ENGINE_TYPE_UNSTRUCTURED = "unstructured" +ENGINE_TYPE_STRUCTURED = "structured" + + +class IdDocument(BaseModel): + """Default structured 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") + + +STRUCTURED_OUTPUT_MODEL = IdDocument + + +def structured_ocr_prompt(model: type[BaseModel] = STRUCTURED_OUTPUT_MODEL) -> str: + """Build OCR instruction text from the pydantic structured-output schema.""" + header = (model.__doc__ or "Extract structured 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", "paddle,openai,gemini") + return [engine_id.strip() for engine_id in raw.split(",") if engine_id.strip()] + + +ENGINE_REGISTRY: dict[str, dict[str, Any]] = { + "paddle": { + "type": ENGINE_TYPE_UNSTRUCTURED, + "provider": "paddle", + "model": "PP-OCRv6", + "init_kwargs": { + "text_detection_model_name": "PP-OCRv6_tiny_det", + "text_recognition_model_name": "PP-OCRv6_tiny_rec", + "use_doc_orientation_classify": False, + "use_doc_unwarping": False, + "use_textline_orientation": False, + "engine": "transformers", + }, + }, + "openai": { + "type": ENGINE_TYPE_STRUCTURED, + "provider": "openai", + "model": "gpt-5.4-nano", + "init_kwargs": { + "temperature": 0, + }, + }, + "gemini": { + "type": ENGINE_TYPE_STRUCTURED, + "provider": "gemini", + "model": "gemini-3.5-flash", + "init_kwargs": { + "temperature": 0, + }, + }, +} + +ENABLED_ENGINES = [engine_id for engine_id in _parse_ocr_engines() if engine_id in ENGINE_REGISTRY] + + +def engines_of_type(engine_type: str) -> list[str]: + return [engine_id for engine_id in ENABLED_ENGINES if ENGINE_REGISTRY[engine_id]["type"] == engine_type] + + +def fo_text_field(engine_id: str) -> str: + return f"{engine_id}_ocr" + + +def fo_det_field(engine_id: str) -> str: + return f"{engine_id}_boxes" + + +def is_unstructured_engine(engine_id: str) -> bool: + return ENGINE_REGISTRY[engine_id]["type"] == ENGINE_TYPE_UNSTRUCTURED + + +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_unstructured_images() -> bool: + return _local_dir_has_images(LOCAL_UNSTRUCTURED_DIR) + + +def use_local_structured_images() -> bool: + return _local_dir_has_images(LOCAL_STRUCTURED_DIR) diff --git a/examples/ocr/data.py b/examples/ocr/data.py new file mode 100644 index 00000000..521b4537 --- /dev/null +++ b/examples/ocr/data.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +from cv_pipeliner.utils.fiftyone import FifyOneSession +from datapipe.compute import Table +from datapipe.store.database import TableStoreDB +from datapipe.store.filedir import PILFile, TableStoreFiledir +from datapipe_ml.utils.image_data_stores import FiftyOneImagesDataTableStore +from sqlalchemy import JSON, Column, String, Text + +from config import ( + COMPOSITES_DIR, + DBCONN, + ENABLED_ENGINES, + ENGINE_REGISTRY, + ENGINE_TYPE_UNSTRUCTURED, + FIFTYONE_DATASET_NAME, + fo_det_field, + fo_text_field, + is_unstructured_engine, +) + +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("kind", String(64)), + 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("engine_type", String(64)), + 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("engine_type", String(64)), + Column("boxes", JSON), + Column("texts", JSON), + Column("scores", JSON), + Column("structured_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, + fo_detections_label=fo_det_field(engine_id) if is_unstructured_engine(engine_id) else None, + additional_info_keys_in_sample=[fo_text_field(engine_id)], + additional_info_keys_in_fo_detections=["engine_id"], + rm_only_fo_fields=True, + primary_schema=[Column("image_id", String(255), primary_key=True)], + ), + ) + +composites_tbl = Table( + name="composites", + store=TableStoreFiledir( + str(COMPOSITES_DIR / "{engine_id}" / "{image_id}.jpg"), + PILFile("jpg"), + primary_schema=[ + Column("engine_id", String(64), primary_key=True), + Column("image_id", String(255), primary_key=True), + ], + ), +) + +UNSTRUCTURED_ENGINE_IDS = [ + engine_id for engine_id in ENABLED_ENGINES if ENGINE_REGISTRY[engine_id]["type"] == ENGINE_TYPE_UNSTRUCTURED +] diff --git a/examples/ocr/engines.py b/examples/ocr/engines.py new file mode 100644 index 00000000..67b40e4c --- /dev/null +++ b/examples/ocr/engines.py @@ -0,0 +1,195 @@ +from __future__ import annotations + +import base64 +import json +import logging +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from config import ( + ENGINE_REGISTRY, + GEMINI_API_KEY, + IdDocument, + OPENAI_API_KEY, + STRUCTURED_OUTPUT_MODEL, + structured_ocr_prompt, +) + +logger = logging.getLogger(__name__) + +_paddle_ocr_instance = None + + +@dataclass +class OcrResult: + has_boxes: bool + boxes: list[list[float]] + texts: list[str] + scores: list[float] + structured_json: str | None + full_text: str | None + + +def _get_paddle_ocr(): + global _paddle_ocr_instance + if _paddle_ocr_instance is None: + from paddleocr import PaddleOCR + + init_kwargs = ENGINE_REGISTRY["paddle"]["init_kwargs"] + _paddle_ocr_instance = PaddleOCR(**init_kwargs) + return _paddle_ocr_instance + + +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 run_paddle(image_path: str, engine_id: str = "paddle", **_kwargs: Any) -> OcrResult: + del engine_id + output = _get_paddle_ocr().predict(input=image_path)[0] + boxes = [list(map(float, box)) for box in output.get("rec_boxes", [])] + texts = [str(text) for text in output.get("rec_texts", [])] + scores = [float(score) for score in output.get("rec_scores", [])] + full_text = "\n".join(texts) + return OcrResult( + has_boxes=len(boxes) > 0, + boxes=boxes, + texts=texts, + scores=scores, + structured_json=None, + full_text=full_text, + ) + + +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 = structured_ocr_prompt(STRUCTURED_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=STRUCTURED_OUTPUT_MODEL, + **init_kwargs, + ) + document = response.output_parsed + if document is None: + raise ValueError("OpenAI structured parse returned empty output") + if not isinstance(document, IdDocument): + document = IdDocument.model_validate(document) + + structured_json = document.model_dump_json(indent=2) + return OcrResult( + has_boxes=False, + boxes=[], + texts=[], + scores=[], + structured_json=structured_json, + full_text=structured_json, + ) + + +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 = structured_ocr_prompt(STRUCTURED_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=STRUCTURED_OUTPUT_MODEL, + **init_kwargs, + ), + ) + document = getattr(response, "parsed", None) + if document is None: + document = IdDocument.model_validate_json(response.text or "{}") + elif not isinstance(document, IdDocument): + document = IdDocument.model_validate(document) + + structured_json = document.model_dump_json(indent=2) + return OcrResult( + has_boxes=False, + boxes=[], + texts=[], + scores=[], + structured_json=structured_json, + full_text=structured_json, + ) + + +ENGINE_RUNNERS = { + "paddle": run_paddle, + "openai": run_openai, + "gemini": run_gemini, +} + + +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) + + +def serialize_boxes(boxes: list[list[float]]) -> str: + return json.dumps(boxes) + + +def serialize_texts(texts: list[str]) -> str: + return json.dumps(texts) + + +def serialize_scores(scores: list[float]) -> str: + return json.dumps(scores) diff --git a/examples/ocr/pyproject.toml b/examples/ocr/pyproject.toml new file mode 100644 index 00000000..9a01ed6e --- /dev/null +++ b/examples/ocr/pyproject.toml @@ -0,0 +1,30 @@ +[project] +name = "ocr-example" +version = "0" +requires-python = ">=3.10,<3.13" +dependencies = [ + "datapipe-ml[torch,fiftyone]", + "transformers==5.12.1", + "python-dotenv>=1.0", + "torch==2.6.0", + "torchvision==0.21.0", + "stringzilla==4.4.0", # needed for albumentations + "fiftyone==1.17.0", + "paddleocr==3.7.0", + "openai>=1.0", + "google-genai>=1.0", + "datasets>=3.0", + "pydantic>=2.0", +] + +[[tool.uv.index]] +name = "pytorch-cu124" +url = "https://download.pytorch.org/whl/cu124" +explicit = true + +[tool.uv.sources] +torch = { index = "pytorch-cu124" } +torchvision = { index = "pytorch-cu124" } +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" } diff --git a/examples/ocr/steps.py b/examples/ocr/steps.py new file mode 100644 index 00000000..aeb77001 --- /dev/null +++ b/examples/ocr/steps.py @@ -0,0 +1,343 @@ +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 BboxData, ImageData +from PIL import Image + +from config import ( + DATA_DIR, + ENGINE_REGISTRY, + ENGINE_TYPE_STRUCTURED, + ENGINE_TYPE_UNSTRUCTURED, + ENABLED_ENGINES, + HF_DATASET_STRUCTURED, + HF_DATASET_UNSTRUCTURED, + HF_LIMIT_STRUCTURED, + HF_LIMIT_UNSTRUCTURED, + HF_SPLIT_STRUCTURED, + HF_SPLIT_UNSTRUCTURED, + IMAGE_SUFFIXES, + LOCAL_STRUCTURED_DIR, + LOCAL_UNSTRUCTURED_DIR, + fo_text_field, + is_unstructured_engine, + use_local_structured_images, + use_local_unstructured_images, +) +from engines import run_engine +from viz import visualize_ocr_side_by_side + +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 _image_id(kind: str, stem: str) -> str: + return f"{kind}__{stem}" + + +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(kind: str, 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": _image_id(kind, path.stem), + "kind": kind, + "image_path": str(path.resolve()), + } + for path in files + ] + + +def _materialize_unstructured_dataset() -> list[dict[str, str]]: + if use_local_unstructured_images(): + assert LOCAL_UNSTRUCTURED_DIR is not None + records = _list_images_from_local_dir(ENGINE_TYPE_UNSTRUCTURED, LOCAL_UNSTRUCTURED_DIR) + logger.info("Using local unstructured images from %s: %d images", LOCAL_UNSTRUCTURED_DIR, len(records)) + return records + + ds = _load_hf_split(HF_DATASET_UNSTRUCTURED, HF_SPLIT_UNSTRUCTURED) + save_dir = _dataset_cache_dir(HF_DATASET_UNSTRUCTURED) + records: list[dict[str, str]] = [] + + for idx, item in enumerate(ds): + if idx >= HF_LIMIT_UNSTRUCTURED: + break + + filename = item.get("filename") + if filename: + stem = Path(str(filename)).stem + out_path = save_dir / str(filename) + else: + stem = f"item_{idx:06d}" + out_path = save_dir / f"{stem}.jpg" + + if not out_path.exists(): + image = item.get("image") + if image is None: + logger.warning("Skipping unstructured sample %s: no image field", idx) + continue + _save_pil_image(image, out_path) + + records.append( + { + "image_id": _image_id(ENGINE_TYPE_UNSTRUCTURED, stem), + "kind": ENGINE_TYPE_UNSTRUCTURED, + "image_path": str(out_path.resolve()), + } + ) + + logger.info( + "Using HF unstructured dataset %s: %d images", + HF_DATASET_UNSTRUCTURED, + len(records), + ) + return records + + +def _structured_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_structured_dataset() -> list[dict[str, str]]: + if use_local_structured_images(): + assert LOCAL_STRUCTURED_DIR is not None + records = _list_images_from_local_dir(ENGINE_TYPE_STRUCTURED, LOCAL_STRUCTURED_DIR) + logger.info("Using local structured images from %s: %d images", LOCAL_STRUCTURED_DIR, len(records)) + return records + + ds = _load_hf_split(HF_DATASET_STRUCTURED, HF_SPLIT_STRUCTURED) + save_dir = _dataset_cache_dir(HF_DATASET_STRUCTURED) + records: list[dict[str, str]] = [] + + for idx, item in enumerate(ds): + if idx >= HF_LIMIT_STRUCTURED: + break + + stem = _structured_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 structured sample %s: no image field", idx) + continue + _save_pil_image(image, out_path) + + records.append( + { + "image_id": _image_id(ENGINE_TYPE_STRUCTURED, stem), + "kind": ENGINE_TYPE_STRUCTURED, + "image_path": str(out_path.resolve()), + } + ) + + logger.info( + "Using HF structured dataset %s: %d images", + HF_DATASET_STRUCTURED, + len(records), + ) + return records + + +def list_images() -> Iterator[pd.DataFrame]: + records = _materialize_unstructured_dataset() + _materialize_structured_dataset() + if not records: + raise ValueError( + "No images found. Set LOCAL_UNSTRUCTURED_DIR / LOCAL_STRUCTURED_DIR with images " + "or configure HF datasets in .env" + ) + yield pd.DataFrame(records, columns=["image_id", "kind", "image_path"]) + + +def list_engines() -> Iterator[pd.DataFrame]: + rows = [ + { + "engine_id": engine_id, + "engine_type": ENGINE_REGISTRY[engine_id]["type"], + "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", + "engine_type", + "boxes", + "texts", + "scores", + "structured_json", + "full_text", + ] + ) + + merged = images_df.merge(engines_df, how="cross") + merged = merged[merged["kind"] == merged["engine_type"]] + if merged.empty: + return pd.DataFrame( + columns=[ + "image_id", + "engine_id", + "engine_type", + "boxes", + "texts", + "scores", + "structured_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"], + "engine_type": row["engine_type"], + "boxes": result.boxes, + "texts": result.texts, + "scores": result.scores, + "structured_json": result.structured_json, + "full_text": result.full_text, + } + ) + + return pd.DataFrame(records) + + +def _build_unstructured_image_data(image_path: str, engine_id: str, ocr_row: pd.Series) -> ImageData: + boxes = ocr_row.get("boxes") or [] + texts = ocr_row.get("texts") or [] + scores = ocr_row.get("scores") or [] + text_field = fo_text_field(engine_id) + + bboxes_data = [] + for box, text, score in zip(boxes, texts, scores): + if len(box) == 4: + xmin, ymin, xmax, ymax = [float(v) for v in box] + else: + xs = [float(p[0]) for p in box] + ys = [float(p[1]) for p in box] + xmin, ymin, xmax, ymax = min(xs), min(ys), max(xs), max(ys) + + bboxes_data.append( + BboxData( + xmin=xmin, + ymin=ymin, + xmax=xmax, + ymax=ymax, + label=str(text), + detection_score=float(score), + additional_info={"engine_id": engine_id}, + ) + ) + + full_text = ocr_row.get("full_text") + if full_text is None and texts: + full_text = "\n".join(str(text) for text in texts) + + return ImageData( + image_path=image_path, + bboxes_data=bboxes_data, + additional_info={text_field: full_text or ""}, + ) + + +def _build_structured_image_data(image_path: str, engine_id: str, ocr_row: pd.Series) -> ImageData: + text_field = fo_text_field(engine_id) + structured_json = ocr_row.get("structured_json") or ocr_row.get("full_text") or "" + if isinstance(structured_json, dict): + structured_json = json.dumps(structured_json, indent=2) + return ImageData( + image_path=image_path, + bboxes_data=[], + additional_info={text_field: str(structured_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_path = str(row["image_path"]) + if is_unstructured_engine(engine_id): + image_data = _build_unstructured_image_data(image_path, engine_id, row) + else: + image_data = _build_structured_image_data(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 + + +def render_composite(images_df: pd.DataFrame, ocr_results_df: pd.DataFrame) -> pd.DataFrame: + if images_df.empty or ocr_results_df.empty: + return pd.DataFrame(columns=["engine_id", "image_id", "image"]) + + unstructured_results = ocr_results_df[ocr_results_df["engine_type"] == ENGINE_TYPE_UNSTRUCTURED] + if unstructured_results.empty: + return pd.DataFrame(columns=["engine_id", "image_id", "image"]) + + merged = images_df.merge(unstructured_results, on="image_id", how="inner") + records = [] + for _, row in merged.iterrows(): + image = Image.open(row["image_path"]).convert("RGB") + composite = visualize_ocr_side_by_side( + image=image, + boxes=row.get("boxes") or [], + texts=row.get("texts") or [], + scores=row.get("scores") or [], + ) + records.append( + { + "engine_id": row["engine_id"], + "image_id": row["image_id"], + "image": composite, + } + ) + + return pd.DataFrame(records, columns=["engine_id", "image_id", "image"]) diff --git a/examples/ocr/viz.py b/examples/ocr/viz.py new file mode 100644 index 00000000..e874a053 --- /dev/null +++ b/examples/ocr/viz.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +from typing import Iterable, Sequence + +from PIL import Image, ImageDraw, ImageFont + + +def _fit_font_to_box( + draw: ImageDraw.ImageDraw, + text: str, + box_w: float, + box_h: float, + max_size: int = 64, + min_size: int = 8, +) -> ImageFont.ImageFont: + """Return largest font that fits text into target box.""" + box_w = max(1.0, box_w) + box_h = max(1.0, box_h) + for size in range(max_size, min_size - 1, -1): + font = ImageFont.load_default(size) + left, top, right, bottom = draw.textbbox((0, 0), text, font=font) + text_w = right - left + text_h = bottom - top + if text_w <= box_w and text_h <= box_h: + return font + return ImageFont.load_default(min_size) + + +def visualize_ocr_side_by_side( + image: Image.Image, + boxes: Iterable[Sequence[Sequence[float] | float]], + texts: Iterable[str], + scores: Iterable[float], +) -> Image.Image: + """Return image with right-side OCR visualization panel. + + Output image size: (2 * width, height). + Left side: original image. + Right side: white canvas with OCR boxes and labels. + """ + width, height = image.size + result = Image.new("RGB", (width * 2, height), "white") + result.paste(image.convert("RGB"), (0, 0)) + + draw = ImageDraw.Draw(result) + x_offset = width + + for box, text, score in zip(boxes, texts, scores): + if len(box) == 4 and isinstance(box[0], (int, float)): + x1, y1, x2, y2 = [float(v) for v in box] # type: ignore[arg-type] + polygon = [(x1, y1), (x2, y1), (x2, y2), (x1, y2)] + else: + polygon = [(float(p[0]), float(p[1])) for p in box] # type: ignore[index] + + draw.polygon(polygon, outline="red", width=2) + shifted = [(x + x_offset, y) for x, y in polygon] + draw.polygon(shifted, outline="red", width=2) + + min_x = min(x for x, _ in shifted) + max_x = max(x for x, _ in shifted) + min_y = min(y for _, y in shifted) + max_y = max(y for _, y in shifted) + box_w = max_x - min_x + box_h = max_y - min_y + + main_font = _fit_font_to_box(draw, text, box_w * 0.9, box_h * 0.55) + main_left, main_top, main_right, main_bottom = draw.textbbox((0, 0), text, font=main_font) + main_w = main_right - main_left + main_h = main_bottom - main_top + text_x = min_x + (box_w - main_w) / 2 + text_y = min_y + (box_h - main_h) / 2 + draw.text((text_x, text_y), text, fill="blue", font=main_font) + + conf_text = f"{float(score):.3f}" + conf_font = _fit_font_to_box(draw, conf_text, box_w * 0.45, box_h * 0.25, max_size=28, min_size=6) + draw.text((min_x + 2, min_y + 2), conf_text, fill="green", font=conf_font) + + return result From 9fa389731219a7f079b32efaefa04bed7d997f13 Mon Sep 17 00:00:00 2001 From: korotas Date: Fri, 3 Jul 2026 15:05:18 +0200 Subject: [PATCH 2/6] remove unstructured; add qwen; remove torch gpu dep --- examples/ocr/.env.example | 19 ++- examples/ocr/README.md | 60 ++++----- examples/ocr/app.py | 12 +- examples/ocr/config.py | 82 ++++--------- examples/ocr/data.py | 41 +------ examples/ocr/engines.py | 137 +++++++++------------ examples/ocr/pyproject.toml | 14 +-- examples/ocr/steps.py | 238 ++++++------------------------------ examples/ocr/viz.py | 78 ------------ 9 files changed, 164 insertions(+), 517 deletions(-) delete mode 100644 examples/ocr/viz.py diff --git a/examples/ocr/.env.example b/examples/ocr/.env.example index e7ce50cf..17ebfc54 100644 --- a/examples/ocr/.env.example +++ b/examples/ocr/.env.example @@ -1,17 +1,12 @@ DB_URL=postgresql+psycopg2://postgres:postgres@localhost:5432/postgres -OCR_ENGINES=paddle,openai,gemini +OCR_ENGINES=openai,gemini,qwen OPENAI_API_KEY=replace-me GEMINI_API_KEY=replace-me -# local mode: use when dir exists and contains images (optional HF fallback per branch) -# LOCAL_UNSTRUCTURED_DIR=./data/unstructured -# LOCAL_STRUCTURED_DIR=./data/structured -# unstructured branch (paddle): free-form document scans -HF_DATASET_UNSTRUCTURED=MLap/Book-Scan-OCR -HF_SPLIT_UNSTRUCTURED=train -# structured branch (openai/gemini): id documents -> IdDocument(name,surname,birthdate) -HF_DATASET_STRUCTURED=ud-synthetic/printed-usa-passports -HF_SPLIT_STRUCTURED=train -HF_LIMIT_STRUCTURED=10 -HF_LIMIT_UNSTRUCTURED=10 +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/README.md b/examples/ocr/README.md index 5f0124ce..883d7341 100644 --- a/examples/ocr/README.md +++ b/examples/ocr/README.md @@ -1,43 +1,32 @@ # OCR + FiftyOne Pipeline Datapipe example for: -- ingesting two Hugging Face datasets (unstructured book scans + structured passport images) -- running OCR with a config-driven engine switch (PaddleOCR v6, OpenAI, Gemini) -- publishing results to FiftyOne (native boxes for unstructured, JSON StringFields for structured) -- optional side-by-side composite renders for unstructured OCR +- 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 -Each branch supports **local folder** or **Hugging Face dataset** fallback. +Supports **local folder** or **Hugging Face dataset** fallback. -### Unstructured (`paddle`) - -- Local: `LOCAL_UNSTRUCTURED_DIR` with `.jpg`/`.jpeg`/`.png`/`.webp`/`.heic` inside (recursive scan) -- HF fallback when local dir unset, missing, or empty: `HF_DATASET_UNSTRUCTURED` (default `MLap/Book-Scan-OCR`) +- 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//` -- Routed to engines with `engine_type=unstructured` - -### Structured (`openai`, `gemini`) -- Local: `LOCAL_STRUCTURED_DIR` with supported images inside (recursive scan) -- HF fallback when local dir unset, missing, or empty: `HF_DATASET_STRUCTURED` (default `ud-synthetic/printed-usa-passports`) -- HF images cached under `DATA_DIR//` -- Routed to engines with `engine_type=structured` -- Both LLMs share the same pydantic schema `IdDocument(name, surname, birthdate)` from `config.py` - -Use `HF_LIMIT_UNSTRUCTURED` / `HF_LIMIT_STRUCTURED` to cap HF download size and LLM cost (ignored for local dirs). +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=paddle,openai,gemini +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 (PaddleOCR init, LLM `temperature`, etc.) live in `ENGINE_REGISTRY` inside `config.py`, not in `.env`. +Model names and inference kwargs (LLM `temperature`, etc.) live in `ENGINE_REGISTRY` inside `config.py`, not in `.env`. ## Prerequisites @@ -49,16 +38,17 @@ Set `DB_URL` in `.env`: postgresql+psycopg2://user:password@host:port/dbname ``` -Create tables: +Create tables (use a fresh DB or drop old tables after schema changes): ```shell datapipe db create-all ``` -### API keys (structured engines) +### API keys - `OPENAI_API_KEY` for `openai` - `GEMINI_API_KEY` (or `GOOGLE_API_KEY`) for `gemini` +- `QWEN_API_KEY` for `qwen` (DashScope compatible OpenAI API) ### FiftyOne App @@ -68,7 +58,7 @@ Launch on the same machine as the pipeline: fiftyone app launch --remote --address 0.0.0.0 --port 5151 --wait -1 ``` -### Caption Viewer plugin (structured JSON) +### Caption Viewer plugin Install the community plugin for readable JSON in the FiftyOne modal: @@ -76,12 +66,12 @@ Install the community plugin for readable JSON in the FiftyOne modal: fiftyone plugins download https://github.com/harpreetsahota204/caption_viewer ``` -In the App: open a sample → add panel → **Caption Viewer** → select `openai_ocr` or `gemini_ocr`. +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, datasets, limits) +2. Edit `.env` (DB, API keys, dataset, limits) 3. Install deps: `uv sync` 4. From this folder: @@ -95,22 +85,20 @@ Or run stages by label: datapipe step --labels=stage=ingest run datapipe step --labels=stage=ocr run datapipe step --labels=stage=fiftyone run -datapipe step --labels=stage=composite run ``` ## Visualize in FiftyOne Dataset name: `FIFTYONE_DATASET_NAME` (default `datapipe_ocr`). -| Engine type | FiftyOne fields | How to view | -|-------------|-----------------|-------------| -| unstructured (`paddle`) | `paddle_boxes` detections, `paddle_ocr` text | Toggle detections overlay; box label = recognized text | -| structured (`openai`, `gemini`) | `openai_ocr`, `gemini_ocr` StringFields | Caption Viewer panel (pretty JSON) | - -Open multiple Caption Viewer panels to compare `openai_ocr` vs `gemini_ocr` side by side. +| 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) | -Composite PNGs (unstructured only) are written to `data/composites/{engine_id}/` when running `stage=composite`. +Open multiple Caption Viewer panels to compare engine fields side by side. -## Customize structured schema +## Customize output schema -Edit `IdDocument` in `config.py`. Both OpenAI and Gemini use it as structured output schema. The OCR instruction text is generated automatically from that model via `structured_ocr_prompt()`. +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 index cfbee28f..d27f3c8f 100644 --- a/examples/ocr/app.py +++ b/examples/ocr/app.py @@ -32,7 +32,7 @@ inputs=[data.images_tbl, data.engines_tbl], outputs=[data.ocr_results_tbl], transform_keys=["image_id", "engine_id"], - chunk_size=8, + chunk_size=2, labels=[("stage", "ocr")], ), ] @@ -48,16 +48,6 @@ ) ) -pipeline_steps.append( - BatchTransform( - func=steps.render_composite, - inputs=[data.images_tbl, data.ocr_results_tbl], - outputs=[data.composites_tbl], - transform_keys=["engine_id", "image_id"], - labels=[("stage", "composite")], - ) -) - pipeline = Pipeline(pipeline_steps) ds = DataStore(DBCONN, create_meta_table=True) diff --git a/examples/ocr/config.py b/examples/ocr/config.py index b99b3ea2..4d47d194 100644 --- a/examples/ocr/config.py +++ b/examples/ocr/config.py @@ -9,51 +9,40 @@ from pydantic import BaseModel, Field DATA_DIR = Path(os.environ.get("DATA_DIR", "./data")).resolve() -COMPOSITES_DIR = DATA_DIR / "composites" -_local_unstructured_dir_raw = os.environ.get("LOCAL_UNSTRUCTURED_DIR") -LOCAL_UNSTRUCTURED_DIR = ( - Path(_local_unstructured_dir_raw).resolve() if _local_unstructured_dir_raw else None -) -_local_structured_dir_raw = os.environ.get("LOCAL_STRUCTURED_DIR") -LOCAL_STRUCTURED_DIR = Path(_local_structured_dir_raw).resolve() if _local_structured_dir_raw else None +_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_UNSTRUCTURED = os.environ.get("HF_DATASET_UNSTRUCTURED", "MLap/Book-Scan-OCR") -HF_SPLIT_UNSTRUCTURED = os.environ.get("HF_SPLIT_UNSTRUCTURED", "train") -HF_LIMIT_UNSTRUCTURED = int(os.environ.get("HF_LIMIT_UNSTRUCTURED", "50")) - -HF_DATASET_STRUCTURED = os.environ.get("HF_DATASET_STRUCTURED", "ud-synthetic/printed-usa-passports") -HF_SPLIT_STRUCTURED = os.environ.get("HF_SPLIT_STRUCTURED", "train") -HF_LIMIT_STRUCTURED = int(os.environ.get("HF_LIMIT_STRUCTURED", "50")) +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") -ENGINE_TYPE_UNSTRUCTURED = "unstructured" -ENGINE_TYPE_STRUCTURED = "structured" - class IdDocument(BaseModel): - """Default structured OCR output for passport / id-doc images.""" + """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") -STRUCTURED_OUTPUT_MODEL = IdDocument +OUTPUT_MODEL = IdDocument -def structured_ocr_prompt(model: type[BaseModel] = STRUCTURED_OUTPUT_MODEL) -> str: - """Build OCR instruction text from the pydantic structured-output schema.""" - header = (model.__doc__ or "Extract structured fields from the document image.").strip() +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() @@ -62,26 +51,12 @@ def structured_ocr_prompt(model: type[BaseModel] = STRUCTURED_OUTPUT_MODEL) -> s def _parse_ocr_engines() -> list[str]: - raw = os.environ.get("OCR_ENGINES", "paddle,openai,gemini") + 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]] = { - "paddle": { - "type": ENGINE_TYPE_UNSTRUCTURED, - "provider": "paddle", - "model": "PP-OCRv6", - "init_kwargs": { - "text_detection_model_name": "PP-OCRv6_tiny_det", - "text_recognition_model_name": "PP-OCRv6_tiny_rec", - "use_doc_orientation_classify": False, - "use_doc_unwarping": False, - "use_textline_orientation": False, - "engine": "transformers", - }, - }, "openai": { - "type": ENGINE_TYPE_STRUCTURED, "provider": "openai", "model": "gpt-5.4-nano", "init_kwargs": { @@ -89,34 +64,33 @@ def _parse_ocr_engines() -> list[str]: }, }, "gemini": { - "type": ENGINE_TYPE_STRUCTURED, "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 engines_of_type(engine_type: str) -> list[str]: - return [engine_id for engine_id in ENABLED_ENGINES if ENGINE_REGISTRY[engine_id]["type"] == engine_type] - - def fo_text_field(engine_id: str) -> str: return f"{engine_id}_ocr" -def fo_det_field(engine_id: str) -> str: - return f"{engine_id}_boxes" - - -def is_unstructured_engine(engine_id: str) -> bool: - return ENGINE_REGISTRY[engine_id]["type"] == ENGINE_TYPE_UNSTRUCTURED - - def _local_dir_has_images(local_dir: Path | None) -> bool: if local_dir is None or not local_dir.exists(): return False @@ -125,9 +99,5 @@ def _local_dir_has_images(local_dir: Path | None) -> bool: ) -def use_local_unstructured_images() -> bool: - return _local_dir_has_images(LOCAL_UNSTRUCTURED_DIR) - - -def use_local_structured_images() -> bool: - return _local_dir_has_images(LOCAL_STRUCTURED_DIR) +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 index 521b4537..5a36dd55 100644 --- a/examples/ocr/data.py +++ b/examples/ocr/data.py @@ -3,21 +3,10 @@ from cv_pipeliner.utils.fiftyone import FifyOneSession from datapipe.compute import Table from datapipe.store.database import TableStoreDB -from datapipe.store.filedir import PILFile, TableStoreFiledir from datapipe_ml.utils.image_data_stores import FiftyOneImagesDataTableStore -from sqlalchemy import JSON, Column, String, Text +from sqlalchemy import Column, String, Text -from config import ( - COMPOSITES_DIR, - DBCONN, - ENABLED_ENGINES, - ENGINE_REGISTRY, - ENGINE_TYPE_UNSTRUCTURED, - FIFTYONE_DATASET_NAME, - fo_det_field, - fo_text_field, - is_unstructured_engine, -) +from config import DBCONN, ENABLED_ENGINES, FIFTYONE_DATASET_NAME, fo_text_field fo_session = FifyOneSession() @@ -28,7 +17,6 @@ name="images", data_sql_schema=[ Column("image_id", String(255), primary_key=True), - Column("kind", String(64)), Column("image_path", String(2048)), ], create_table=True, @@ -42,7 +30,6 @@ name="engines", data_sql_schema=[ Column("engine_id", String(64), primary_key=True), - Column("engine_type", String(64)), Column("provider", String(64)), Column("model", String(255)), ], @@ -58,11 +45,7 @@ data_sql_schema=[ Column("image_id", String(255), primary_key=True), Column("engine_id", String(64), primary_key=True), - Column("engine_type", String(64)), - Column("boxes", JSON), - Column("texts", JSON), - Column("scores", JSON), - Column("structured_json", Text), + Column("output_json", Text), Column("full_text", Text), ], create_table=True, @@ -76,26 +59,8 @@ store=FiftyOneImagesDataTableStore( dataset=FIFTYONE_DATASET_NAME, fo_session=fo_session, - fo_detections_label=fo_det_field(engine_id) if is_unstructured_engine(engine_id) else None, additional_info_keys_in_sample=[fo_text_field(engine_id)], - additional_info_keys_in_fo_detections=["engine_id"], rm_only_fo_fields=True, primary_schema=[Column("image_id", String(255), primary_key=True)], ), ) - -composites_tbl = Table( - name="composites", - store=TableStoreFiledir( - str(COMPOSITES_DIR / "{engine_id}" / "{image_id}.jpg"), - PILFile("jpg"), - primary_schema=[ - Column("engine_id", String(64), primary_key=True), - Column("image_id", String(255), primary_key=True), - ], - ), -) - -UNSTRUCTURED_ENGINE_IDS = [ - engine_id for engine_id in ENABLED_ENGINES if ENGINE_REGISTRY[engine_id]["type"] == ENGINE_TYPE_UNSTRUCTURED -] diff --git a/examples/ocr/engines.py b/examples/ocr/engines.py index 67b40e4c..7379045d 100644 --- a/examples/ocr/engines.py +++ b/examples/ocr/engines.py @@ -1,7 +1,6 @@ from __future__ import annotations import base64 -import json import logging from dataclasses import dataclass from pathlib import Path @@ -10,35 +9,19 @@ from config import ( ENGINE_REGISTRY, GEMINI_API_KEY, - IdDocument, OPENAI_API_KEY, - STRUCTURED_OUTPUT_MODEL, - structured_ocr_prompt, + OUTPUT_MODEL, + QWEN_API_KEY, + ocr_prompt, ) logger = logging.getLogger(__name__) -_paddle_ocr_instance = None - @dataclass class OcrResult: - has_boxes: bool - boxes: list[list[float]] - texts: list[str] - scores: list[float] - structured_json: str | None - full_text: str | None - - -def _get_paddle_ocr(): - global _paddle_ocr_instance - if _paddle_ocr_instance is None: - from paddleocr import PaddleOCR - - init_kwargs = ENGINE_REGISTRY["paddle"]["init_kwargs"] - _paddle_ocr_instance = PaddleOCR(**init_kwargs) - return _paddle_ocr_instance + output_json: str + full_text: str def _encode_image_base64(image_path: str) -> str: @@ -62,21 +45,16 @@ def _engine_config(engine_id: str) -> dict[str, Any]: return ENGINE_REGISTRY[engine_id] -def run_paddle(image_path: str, engine_id: str = "paddle", **_kwargs: Any) -> OcrResult: - del engine_id - output = _get_paddle_ocr().predict(input=image_path)[0] - boxes = [list(map(float, box)) for box in output.get("rec_boxes", [])] - texts = [str(text) for text in output.get("rec_texts", [])] - scores = [float(score) for score in output.get("rec_scores", [])] - full_text = "\n".join(texts) - return OcrResult( - has_boxes=len(boxes) > 0, - boxes=boxes, - texts=texts, - scores=scores, - structured_json=None, - full_text=full_text, - ) +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: @@ -87,7 +65,7 @@ def run_openai(image_path: str, engine_id: str = "openai", **_kwargs: Any) -> Oc engine_cfg = _engine_config(engine_id) init_kwargs = engine_cfg.get("init_kwargs", {}) - prompt = structured_ocr_prompt(STRUCTURED_OUTPUT_MODEL) + prompt = ocr_prompt(OUTPUT_MODEL) client = OpenAI(api_key=OPENAI_API_KEY) mime = _guess_mime_type(image_path) @@ -106,24 +84,13 @@ def run_openai(image_path: str, engine_id: str = "openai", **_kwargs: Any) -> Oc ], } ], - text_format=STRUCTURED_OUTPUT_MODEL, + text_format=OUTPUT_MODEL, **init_kwargs, ) document = response.output_parsed if document is None: raise ValueError("OpenAI structured parse returned empty output") - if not isinstance(document, IdDocument): - document = IdDocument.model_validate(document) - - structured_json = document.model_dump_json(indent=2) - return OcrResult( - has_boxes=False, - boxes=[], - texts=[], - scores=[], - structured_json=structured_json, - full_text=structured_json, - ) + return _ocr_result_from_document(document) def run_gemini(image_path: str, engine_id: str = "gemini", **_kwargs: Any) -> OcrResult: @@ -135,7 +102,7 @@ def run_gemini(image_path: str, engine_id: str = "gemini", **_kwargs: Any) -> Oc engine_cfg = _engine_config(engine_id) init_kwargs = engine_cfg.get("init_kwargs", {}) - prompt = structured_ocr_prompt(STRUCTURED_OUTPUT_MODEL) + prompt = ocr_prompt(OUTPUT_MODEL) client = genai.Client(api_key=GEMINI_API_KEY) image_bytes = Path(image_path).read_bytes() @@ -148,31 +115,57 @@ def run_gemini(image_path: str, engine_id: str = "gemini", **_kwargs: Any) -> Oc ], config=types.GenerateContentConfig( response_mime_type="application/json", - response_schema=STRUCTURED_OUTPUT_MODEL, + response_schema=OUTPUT_MODEL, **init_kwargs, ), ) document = getattr(response, "parsed", None) if document is None: - document = IdDocument.model_validate_json(response.text or "{}") - elif not isinstance(document, IdDocument): - document = IdDocument.model_validate(document) - - structured_json = document.model_dump_json(indent=2) - return OcrResult( - has_boxes=False, - boxes=[], - texts=[], - scores=[], - structured_json=structured_json, - full_text=structured_json, + 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 = { - "paddle": run_paddle, "openai": run_openai, "gemini": run_gemini, + "qwen": run_qwen, } @@ -181,15 +174,3 @@ def run_engine(engine_id: str, image_path: str) -> OcrResult: if runner is None: raise KeyError(f"Unknown OCR engine: {engine_id}") return runner(image_path=image_path, engine_id=engine_id) - - -def serialize_boxes(boxes: list[list[float]]) -> str: - return json.dumps(boxes) - - -def serialize_texts(texts: list[str]) -> str: - return json.dumps(texts) - - -def serialize_scores(scores: list[float]) -> str: - return json.dumps(scores) diff --git a/examples/ocr/pyproject.toml b/examples/ocr/pyproject.toml index 9a01ed6e..13238efc 100644 --- a/examples/ocr/pyproject.toml +++ b/examples/ocr/pyproject.toml @@ -3,28 +3,24 @@ name = "ocr-example" version = "0" requires-python = ">=3.10,<3.13" dependencies = [ - "datapipe-ml[torch,fiftyone]", - "transformers==5.12.1", + "datapipe-ml[fiftyone]", "python-dotenv>=1.0", - "torch==2.6.0", - "torchvision==0.21.0", "stringzilla==4.4.0", # needed for albumentations "fiftyone==1.17.0", - "paddleocr==3.7.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-cu124" -url = "https://download.pytorch.org/whl/cu124" +name = "pytorch-cpu" +url = "https://download.pytorch.org/whl/cpu" explicit = true [tool.uv.sources] -torch = { index = "pytorch-cu124" } -torchvision = { index = "pytorch-cu124" } 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 index aeb77001..dea89ddc 100644 --- a/examples/ocr/steps.py +++ b/examples/ocr/steps.py @@ -6,31 +6,22 @@ from pathlib import Path import pandas as pd -from cv_pipeliner import BboxData, ImageData +from cv_pipeliner import ImageData from PIL import Image from config import ( DATA_DIR, ENGINE_REGISTRY, - ENGINE_TYPE_STRUCTURED, - ENGINE_TYPE_UNSTRUCTURED, ENABLED_ENGINES, - HF_DATASET_STRUCTURED, - HF_DATASET_UNSTRUCTURED, - HF_LIMIT_STRUCTURED, - HF_LIMIT_UNSTRUCTURED, - HF_SPLIT_STRUCTURED, - HF_SPLIT_UNSTRUCTURED, + HF_DATASET_NAME, + HF_LIMIT, + HF_SPLIT, IMAGE_SUFFIXES, - LOCAL_STRUCTURED_DIR, - LOCAL_UNSTRUCTURED_DIR, + LOCAL_IMAGES_DIR, fo_text_field, - is_unstructured_engine, - use_local_structured_images, - use_local_unstructured_images, + use_local_images, ) from engines import run_engine -from viz import visualize_ocr_side_by_side logger = logging.getLogger(__name__) @@ -44,17 +35,13 @@ def _save_pil_image(image: Image.Image, path: Path) -> None: image.convert("RGB").save(path, format="JPEG", quality=95) -def _image_id(kind: str, stem: str) -> str: - return f"{kind}__{stem}" - - 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(kind: str, local_dir: Path) -> list[dict[str, str]]: +def _list_images_from_local_dir(local_dir: Path) -> list[dict[str, str]]: files = sorted( path for path in local_dir.rglob("*") @@ -62,123 +49,74 @@ def _list_images_from_local_dir(kind: str, local_dir: Path) -> list[dict[str, st ) return [ { - "image_id": _image_id(kind, path.stem), - "kind": kind, + "image_id": path.stem, "image_path": str(path.resolve()), } for path in files ] -def _materialize_unstructured_dataset() -> list[dict[str, str]]: - if use_local_unstructured_images(): - assert LOCAL_UNSTRUCTURED_DIR is not None - records = _list_images_from_local_dir(ENGINE_TYPE_UNSTRUCTURED, LOCAL_UNSTRUCTURED_DIR) - logger.info("Using local unstructured images from %s: %d images", LOCAL_UNSTRUCTURED_DIR, len(records)) - return records - - ds = _load_hf_split(HF_DATASET_UNSTRUCTURED, HF_SPLIT_UNSTRUCTURED) - save_dir = _dataset_cache_dir(HF_DATASET_UNSTRUCTURED) - records: list[dict[str, str]] = [] - - for idx, item in enumerate(ds): - if idx >= HF_LIMIT_UNSTRUCTURED: - break - - filename = item.get("filename") - if filename: - stem = Path(str(filename)).stem - out_path = save_dir / str(filename) - else: - stem = f"item_{idx:06d}" - out_path = save_dir / f"{stem}.jpg" - - if not out_path.exists(): - image = item.get("image") - if image is None: - logger.warning("Skipping unstructured sample %s: no image field", idx) - continue - _save_pil_image(image, out_path) - - records.append( - { - "image_id": _image_id(ENGINE_TYPE_UNSTRUCTURED, stem), - "kind": ENGINE_TYPE_UNSTRUCTURED, - "image_path": str(out_path.resolve()), - } - ) - - logger.info( - "Using HF unstructured dataset %s: %d images", - HF_DATASET_UNSTRUCTURED, - len(records), - ) - return records - - -def _structured_sample_stem(item: dict, idx: int) -> str: +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_structured_dataset() -> list[dict[str, str]]: - if use_local_structured_images(): - assert LOCAL_STRUCTURED_DIR is not None - records = _list_images_from_local_dir(ENGINE_TYPE_STRUCTURED, LOCAL_STRUCTURED_DIR) - logger.info("Using local structured images from %s: %d images", LOCAL_STRUCTURED_DIR, len(records)) +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_STRUCTURED, HF_SPLIT_STRUCTURED) - save_dir = _dataset_cache_dir(HF_DATASET_STRUCTURED) + 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_STRUCTURED: + if idx >= HF_LIMIT: break - stem = _structured_sample_stem(item, idx) - out_path = save_dir / f"{stem}.jpg" + 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 structured sample %s: no image field", idx) + logger.warning("Skipping sample %s: no image field", idx) continue _save_pil_image(image, out_path) records.append( { - "image_id": _image_id(ENGINE_TYPE_STRUCTURED, stem), - "kind": ENGINE_TYPE_STRUCTURED, + "image_id": stem, "image_path": str(out_path.resolve()), } ) - logger.info( - "Using HF structured dataset %s: %d images", - HF_DATASET_STRUCTURED, - len(records), - ) + logger.info("Using HF dataset %s: %d images", HF_DATASET_NAME, len(records)) return records def list_images() -> Iterator[pd.DataFrame]: - records = _materialize_unstructured_dataset() + _materialize_structured_dataset() + records = _materialize_images() if not records: raise ValueError( - "No images found. Set LOCAL_UNSTRUCTURED_DIR / LOCAL_STRUCTURED_DIR with images " - "or configure HF datasets in .env" + "No images found. Set LOCAL_IMAGES_DIR with images or configure HF_DATASET_NAME in .env" ) - yield pd.DataFrame(records, columns=["image_id", "kind", "image_path"]) + yield pd.DataFrame(records, columns=["image_id", "image_path"]) def list_engines() -> Iterator[pd.DataFrame]: rows = [ { "engine_id": engine_id, - "engine_type": ENGINE_REGISTRY[engine_id]["type"], "provider": ENGINE_REGISTRY[engine_id]["provider"], "model": ENGINE_REGISTRY[engine_id]["model"], } @@ -191,34 +129,11 @@ def list_engines() -> Iterator[pd.DataFrame]: 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", - "engine_type", - "boxes", - "texts", - "scores", - "structured_json", - "full_text", - ] - ) + return pd.DataFrame(columns=["image_id", "engine_id", "output_json", "full_text"]) merged = images_df.merge(engines_df, how="cross") - merged = merged[merged["kind"] == merged["engine_type"]] if merged.empty: - return pd.DataFrame( - columns=[ - "image_id", - "engine_id", - "engine_type", - "boxes", - "texts", - "scores", - "structured_json", - "full_text", - ] - ) + return pd.DataFrame(columns=["image_id", "engine_id", "output_json", "full_text"]) records = [] for _, row in merged.iterrows(): @@ -227,11 +142,7 @@ def ocr_inference(images_df: pd.DataFrame, engines_df: pd.DataFrame) -> pd.DataF { "image_id": row["image_id"], "engine_id": row["engine_id"], - "engine_type": row["engine_type"], - "boxes": result.boxes, - "texts": result.texts, - "scores": result.scores, - "structured_json": result.structured_json, + "output_json": result.output_json, "full_text": result.full_text, } ) @@ -239,53 +150,15 @@ def ocr_inference(images_df: pd.DataFrame, engines_df: pd.DataFrame) -> pd.DataF return pd.DataFrame(records) -def _build_unstructured_image_data(image_path: str, engine_id: str, ocr_row: pd.Series) -> ImageData: - boxes = ocr_row.get("boxes") or [] - texts = ocr_row.get("texts") or [] - scores = ocr_row.get("scores") or [] +def _build_image_data(image_path: str, engine_id: str, ocr_row: pd.Series) -> ImageData: text_field = fo_text_field(engine_id) - - bboxes_data = [] - for box, text, score in zip(boxes, texts, scores): - if len(box) == 4: - xmin, ymin, xmax, ymax = [float(v) for v in box] - else: - xs = [float(p[0]) for p in box] - ys = [float(p[1]) for p in box] - xmin, ymin, xmax, ymax = min(xs), min(ys), max(xs), max(ys) - - bboxes_data.append( - BboxData( - xmin=xmin, - ymin=ymin, - xmax=xmax, - ymax=ymax, - label=str(text), - detection_score=float(score), - additional_info={"engine_id": engine_id}, - ) - ) - - full_text = ocr_row.get("full_text") - if full_text is None and texts: - full_text = "\n".join(str(text) for text in texts) - - return ImageData( - image_path=image_path, - bboxes_data=bboxes_data, - additional_info={text_field: full_text or ""}, - ) - - -def _build_structured_image_data(image_path: str, engine_id: str, ocr_row: pd.Series) -> ImageData: - text_field = fo_text_field(engine_id) - structured_json = ocr_row.get("structured_json") or ocr_row.get("full_text") or "" - if isinstance(structured_json, dict): - structured_json = json.dumps(structured_json, indent=2) + 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(structured_json)}, + additional_info={text_field: str(output_json)}, ) @@ -301,43 +174,10 @@ def build_image_data(images_df: pd.DataFrame, ocr_results_df: pd.DataFrame) -> p merged = images_df.merge(engine_results, on="image_id", how="inner") records = [] for _, row in merged.iterrows(): - image_path = str(row["image_path"]) - if is_unstructured_engine(engine_id): - image_data = _build_unstructured_image_data(image_path, engine_id, row) - else: - image_data = _build_structured_image_data(image_path, engine_id, row) + 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 - - -def render_composite(images_df: pd.DataFrame, ocr_results_df: pd.DataFrame) -> pd.DataFrame: - if images_df.empty or ocr_results_df.empty: - return pd.DataFrame(columns=["engine_id", "image_id", "image"]) - - unstructured_results = ocr_results_df[ocr_results_df["engine_type"] == ENGINE_TYPE_UNSTRUCTURED] - if unstructured_results.empty: - return pd.DataFrame(columns=["engine_id", "image_id", "image"]) - - merged = images_df.merge(unstructured_results, on="image_id", how="inner") - records = [] - for _, row in merged.iterrows(): - image = Image.open(row["image_path"]).convert("RGB") - composite = visualize_ocr_side_by_side( - image=image, - boxes=row.get("boxes") or [], - texts=row.get("texts") or [], - scores=row.get("scores") or [], - ) - records.append( - { - "engine_id": row["engine_id"], - "image_id": row["image_id"], - "image": composite, - } - ) - - return pd.DataFrame(records, columns=["engine_id", "image_id", "image"]) diff --git a/examples/ocr/viz.py b/examples/ocr/viz.py deleted file mode 100644 index e874a053..00000000 --- a/examples/ocr/viz.py +++ /dev/null @@ -1,78 +0,0 @@ -from __future__ import annotations - -from typing import Iterable, Sequence - -from PIL import Image, ImageDraw, ImageFont - - -def _fit_font_to_box( - draw: ImageDraw.ImageDraw, - text: str, - box_w: float, - box_h: float, - max_size: int = 64, - min_size: int = 8, -) -> ImageFont.ImageFont: - """Return largest font that fits text into target box.""" - box_w = max(1.0, box_w) - box_h = max(1.0, box_h) - for size in range(max_size, min_size - 1, -1): - font = ImageFont.load_default(size) - left, top, right, bottom = draw.textbbox((0, 0), text, font=font) - text_w = right - left - text_h = bottom - top - if text_w <= box_w and text_h <= box_h: - return font - return ImageFont.load_default(min_size) - - -def visualize_ocr_side_by_side( - image: Image.Image, - boxes: Iterable[Sequence[Sequence[float] | float]], - texts: Iterable[str], - scores: Iterable[float], -) -> Image.Image: - """Return image with right-side OCR visualization panel. - - Output image size: (2 * width, height). - Left side: original image. - Right side: white canvas with OCR boxes and labels. - """ - width, height = image.size - result = Image.new("RGB", (width * 2, height), "white") - result.paste(image.convert("RGB"), (0, 0)) - - draw = ImageDraw.Draw(result) - x_offset = width - - for box, text, score in zip(boxes, texts, scores): - if len(box) == 4 and isinstance(box[0], (int, float)): - x1, y1, x2, y2 = [float(v) for v in box] # type: ignore[arg-type] - polygon = [(x1, y1), (x2, y1), (x2, y2), (x1, y2)] - else: - polygon = [(float(p[0]), float(p[1])) for p in box] # type: ignore[index] - - draw.polygon(polygon, outline="red", width=2) - shifted = [(x + x_offset, y) for x, y in polygon] - draw.polygon(shifted, outline="red", width=2) - - min_x = min(x for x, _ in shifted) - max_x = max(x for x, _ in shifted) - min_y = min(y for _, y in shifted) - max_y = max(y for _, y in shifted) - box_w = max_x - min_x - box_h = max_y - min_y - - main_font = _fit_font_to_box(draw, text, box_w * 0.9, box_h * 0.55) - main_left, main_top, main_right, main_bottom = draw.textbbox((0, 0), text, font=main_font) - main_w = main_right - main_left - main_h = main_bottom - main_top - text_x = min_x + (box_w - main_w) / 2 - text_y = min_y + (box_h - main_h) / 2 - draw.text((text_x, text_y), text, fill="blue", font=main_font) - - conf_text = f"{float(score):.3f}" - conf_font = _fit_font_to_box(draw, conf_text, box_w * 0.45, box_h * 0.25, max_size=28, min_size=6) - draw.text((min_x + 2, min_y + 2), conf_text, fill="green", font=conf_font) - - return result From 37b31c3c3d3bd6923717d7572f2731f83be12dac Mon Sep 17 00:00:00 2001 From: korotas Date: Mon, 6 Jul 2026 12:29:11 +0200 Subject: [PATCH 3/6] wip --- examples/ocr/README.md | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/examples/ocr/README.md b/examples/ocr/README.md index 883d7341..83f81207 100644 --- a/examples/ocr/README.md +++ b/examples/ocr/README.md @@ -47,26 +47,24 @@ datapipe db create-all ### API keys - `OPENAI_API_KEY` for `openai` -- `GEMINI_API_KEY` (or `GOOGLE_API_KEY`) for `gemini` +- `GEMINI_API_KEY` for `gemini` - `QWEN_API_KEY` for `qwen` (DashScope compatible OpenAI API) ### FiftyOne App -Launch on the same machine as the pipeline: +First - after creating virtual env install the community plugin for readable JSON in the FiftyOne modal: ```shell -fiftyone app launch --remote --address 0.0.0.0 --port 5151 --wait -1 +fiftyone plugins download https://github.com/harpreetsahota204/caption_viewer ``` -### Caption Viewer plugin - -Install the community plugin for readable JSON in the FiftyOne modal: +Launch on the same machine as the pipeline: ```shell -fiftyone plugins download https://github.com/harpreetsahota204/caption_viewer +fiftyone app launch --remote --address 0.0.0.0 --port 5151 --wait -1 ``` -In the App: open a sample → add panel → **Caption Viewer** → select `openai_ocr`, `gemini_ocr`, or `qwen_ocr`. +To visualize JSON structured output in the App: open a sample → add panel → **Caption Viewer** → select `openai_ocr`, `gemini_ocr`, or `qwen_ocr`. ## Run From 724e7755eb83bda7f4f25fb1d379a11ec0c5c782 Mon Sep 17 00:00:00 2001 From: korotas Date: Mon, 6 Jul 2026 12:42:10 +0200 Subject: [PATCH 4/6] add skill --- .claude/skills/setup-ocr/SKILL.md | 75 +++++++++++++++++++++++++++++++ examples/ocr/README.md | 2 + 2 files changed, 77 insertions(+) create mode 100644 .claude/skills/setup-ocr/SKILL.md diff --git a/.claude/skills/setup-ocr/SKILL.md b/.claude/skills/setup-ocr/SKILL.md new file mode 100644 index 00000000..d5de908f --- /dev/null +++ b/.claude/skills/setup-ocr/SKILL.md @@ -0,0 +1,75 @@ +--- +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)? **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` (models + inference kwargs live in `ENGINE_REGISTRY` in `config.py`, not `.env`). +- **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`. +- **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`). +- **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/README.md b/examples/ocr/README.md index 83f81207..11416e3c 100644 --- a/examples/ocr/README.md +++ b/examples/ocr/README.md @@ -1,5 +1,7 @@ # 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) From 89306dddbc1f34cc846d5d044a7127f69041f2b8 Mon Sep 17 00:00:00 2001 From: korotas Date: Fri, 10 Jul 2026 10:43:47 +0200 Subject: [PATCH 5/6] add ocr to general skill --- .claude/skills/datapipe-examples/SKILL.md | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) 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 From 0c521bce14dc012d8fe4d61ca9b2e2e3ae06292e Mon Sep 17 00:00:00 2001 From: korotas Date: Fri, 10 Jul 2026 10:52:42 +0200 Subject: [PATCH 6/6] upd skill --- .claude/skills/setup-ocr/SKILL.md | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/.claude/skills/setup-ocr/SKILL.md b/.claude/skills/setup-ocr/SKILL.md index d5de908f..bc6e4026 100644 --- a/.claude/skills/setup-ocr/SKILL.md +++ b/.claude/skills/setup-ocr/SKILL.md @@ -9,7 +9,7 @@ description: > 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)? **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? +**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. @@ -17,7 +17,10 @@ This skill = run multi-engine LLM OCR on YOUR passport/id-doc images and publish - **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` (models + inference kwargs live in `ENGINE_REGISTRY` in `config.py`, not `.env`). +- **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. @@ -27,6 +30,8 @@ cross-join, `chunk_size=2`) → `fiftyone` (per-engine `{engine_id}_ocr` StringF - **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` @@ -68,6 +73,7 @@ Open `FIFTYONE_DATASET_NAME`. Per engine: open a sample → add panel → **Capt - **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`.