refactor: solidify startup UX and engine-aware preprocessing
- Slim core deps: move ML stack to optional extras (paddle/tables/figures/scientific/full) - Lazy settings proxy with config search paths (env var, cwd, user dir) - New commands: init, demo, setup [basic|full], first-run guard on run/watch - Engine-aware preprocessing: Paddle gets original image (fixes dark mode 0.83->0.95) - Results table shows Skipped count; lazy run-dir creation - kg_ocr marked experimental with extra, Docker defaults with OCR_PIPELINE_CONFIG - 25/25 tests, ruff clean
This commit is contained in:
21
Dockerfile
21
Dockerfile
@@ -2,25 +2,24 @@ FROM python:3.11-slim
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
tesseract-ocr \
|
||||
libtesseract-dev \
|
||||
poppler-utils \
|
||||
libgl1 \
|
||||
libglib2.0-0 \
|
||||
libsm6 \
|
||||
libxext6 \
|
||||
libxrender-dev \
|
||||
libgomp1 \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
COPY pyproject.toml uv.lock ./
|
||||
RUN pip install --no-cache-dir uv && uv sync --frozen --no-dev
|
||||
RUN pip install --no-cache-dir uv && uv sync --frozen --no-dev --no-install-project
|
||||
|
||||
COPY src ./src
|
||||
COPY config.yaml ./
|
||||
# Optional ML stack: docker build --build-arg EXTRAS="--extra full" .
|
||||
ARG EXTRAS=""
|
||||
RUN uv sync --frozen --no-dev ${EXTRAS}
|
||||
|
||||
ENV PYTHONPATH=/app/src
|
||||
ENV OMP_NUM_THREADS=4
|
||||
COPY docker/config.docker.yaml ./config.docker.yaml
|
||||
|
||||
ENTRYPOINT ["ocr-pipeline"]
|
||||
ENV PATH="/app/.venv/bin:$PATH" \
|
||||
OCR_PIPELINE_CONFIG=/app/config.docker.yaml \
|
||||
OMP_NUM_THREADS=4
|
||||
|
||||
ENTRYPOINT ["ocr-pipeline"]
|
||||
|
||||
82
README.md
82
README.md
@@ -11,61 +11,76 @@ Turns scientific screenshots into RAG-ready Markdown, with metadata attached.
|
||||
|
||||
## Quick start
|
||||
|
||||
The default install is small: Tesseract OCR only, no multi-GB ML downloads.
|
||||
|
||||
```bash
|
||||
uv sync
|
||||
pre-commit install
|
||||
|
||||
# Smoke test a small directory with the lightweight OCR path
|
||||
uv run ocr-pipeline run \
|
||||
--input-dir ~/Pictures/test_screenshots \
|
||||
--output-dir /tmp/ocr-output \
|
||||
--engine tesseract \
|
||||
--workers 1 \
|
||||
--exclude "*.photoslibrary/*"
|
||||
|
||||
uv run ocr-pipeline watch
|
||||
uv run ocr-pipeline reprocess ocr_sc.ipynb
|
||||
uv run ocr-pipeline status
|
||||
uv sync # install the core pipeline
|
||||
uv run ocr-pipeline setup # installs Tesseract if missing (asks first)
|
||||
uv run ocr-pipeline init # write a config for this machine
|
||||
uv run ocr-pipeline demo # verify everything works on a sample screenshot
|
||||
uv run ocr-pipeline run # process your screenshots
|
||||
```
|
||||
|
||||
Config is found in this order: `$OCR_PIPELINE_CONFIG`, `./config.yaml`, then the
|
||||
per-user config directory (`~/Library/Application Support/ocr-pipeline` on macOS,
|
||||
`~/.config/ocr-pipeline` on Linux). Running `run` with neither a config nor
|
||||
`--input-dir` stops with guidance instead of scanning a default directory.
|
||||
|
||||
Use `--force` to reprocess files already recorded in the processing database. The
|
||||
`--exclude` option is repeatable and prevents recursive scans from entering bundles
|
||||
such as macOS `.photoslibrary` directories.
|
||||
|
||||
### Scientific NER environment
|
||||
### Optional capabilities
|
||||
|
||||
`en_core_sci_lg` is distributed separately from scispaCy and is not installed by
|
||||
`uv sync`. The compatible scispaCy 0.5.4 release requires spaCy 3.7.x, so install it
|
||||
in a dedicated Python 3.11 environment rather than mixing it with the main Python
|
||||
3.13 environment:
|
||||
Heavy ML features are opt-in extras and degrade gracefully when absent (a warning
|
||||
is logged and that step is skipped):
|
||||
|
||||
| Extra | Provides |
|
||||
|-------|----------|
|
||||
| `paddle` | PaddleOCR engine (`ocr.engine: paddleocr` or `auto`) |
|
||||
| `scientific` | spaCy/scispaCy NER (`entities.enabled: true`) |
|
||||
| `tables` | Table Transformer detection |
|
||||
| `figures` | LayoutParser figure/caption detection (needs a Detectron2 build) |
|
||||
| `full` | All of the above |
|
||||
|
||||
```bash
|
||||
uv venv --python 3.11 .venv-scispacy
|
||||
source .venv-scispacy/bin/activate
|
||||
uv pip install "spacy>=3.7,<3.8" "scispacy==0.5.4" click
|
||||
uv pip install https://s3-us-west-2.amazonaws.com/ai2-s2-scispacy/releases/v0.5.4/en_core_sci_lg-0.5.4.tar.gz
|
||||
uv pip install -e .
|
||||
uv pip install timm
|
||||
|
||||
# Confirm the model and command use this same environment
|
||||
.venv-scispacy/bin/python -c "import spacy; spacy.load('en_core_sci_lg'); print('scispaCy OK')"
|
||||
.venv-scispacy/bin/ocr-pipeline run --help
|
||||
uv run ocr-pipeline setup full # runs: uv sync --extra full + downloads en_core_sci_lg
|
||||
```
|
||||
|
||||
If the two environments are kept separate, invoke the pipeline with the explicit
|
||||
`.venv-scispacy/bin/ocr-pipeline` path; `uv run` selects the project's default `.venv`.
|
||||
`en_core_sci_lg` is distributed separately from scispaCy. The compatible scispaCy
|
||||
0.5.4 release requires spaCy 3.7.x; if the model install conflicts with your spaCy
|
||||
version, use a dedicated Python 3.11 environment as described in
|
||||
`ocr-pipeline setup full --dry-run` output, and enable `entities` only there.
|
||||
|
||||
## kg_ocr (experimental)
|
||||
|
||||
The top-level `kg_ocr` package is a prototype txtai/litellm RAG interface over the
|
||||
pipeline's output. It is not the supported interface (that is the `ocr-pipeline`
|
||||
CLI) and its dependencies are not installed by default:
|
||||
|
||||
```bash
|
||||
uv sync --extra kg
|
||||
```
|
||||
|
||||
## Docker
|
||||
|
||||
The image ships the Tesseract-only core with a ready-made config that reads from
|
||||
`/data/screenshots` and writes to `/data/output`:
|
||||
|
||||
```bash
|
||||
docker build -t ocr-pipeline .
|
||||
docker run -v ~/Pictures:/data/screenshots -v "$PWD/ocr-output:/data/output" ocr-pipeline run
|
||||
docker run -v ~/Pictures:/data/screenshots -v "$PWD/ocr-output:/data/output" ocr-pipeline watch
|
||||
|
||||
docker run -v ~/Pictures:/data/screenshots -v ./data:/app/data ocr-pipeline run
|
||||
docker run -v ~/Pictures:/data/screenshots -v ./data:/app/data ocr-pipeline watch
|
||||
# ML extras baked in:
|
||||
docker build --build-arg EXTRAS="--extra full" -t ocr-pipeline:full .
|
||||
```
|
||||
|
||||
## Config
|
||||
|
||||
`ocr-pipeline init` writes a minimal working config; the example below shows every
|
||||
knob for reference.
|
||||
|
||||
<details>
|
||||
<summary><code>config.yaml</code></summary>
|
||||
|
||||
@@ -200,6 +215,7 @@ Gene/protein names go through scispaCy's `en_core_sci_lg`. Chemical formulas and
|
||||
|
||||
## Models
|
||||
|
||||
Model weights are only fetched when the matching extra is installed and enabled.
|
||||
First run downloads what it needs, cached in `~/.cache/ocr_pipeline/`:
|
||||
|
||||
- PaddleOCR models (~200MB)
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
# Full-stack config used for this repo's development. New users: run
|
||||
# `ocr-pipeline init` to generate a minimal config instead of copying this file.
|
||||
input:
|
||||
paths:
|
||||
- "~/Pictures"
|
||||
|
||||
15
docker/config.docker.yaml
Normal file
15
docker/config.docker.yaml
Normal file
@@ -0,0 +1,15 @@
|
||||
# Container defaults: host directories are mounted under /data.
|
||||
# Selected via OCR_PIPELINE_CONFIG in the image.
|
||||
input:
|
||||
paths:
|
||||
- "/data/screenshots"
|
||||
|
||||
ocr:
|
||||
engine: "tesseract"
|
||||
|
||||
output:
|
||||
base_directory: "/data/output"
|
||||
write_consolidated: true
|
||||
|
||||
watch:
|
||||
db_path: "/data/output/processed_files.db"
|
||||
@@ -1,5 +1,17 @@
|
||||
from .ocr import get_screenshots, extract_text
|
||||
from .embeddings import create_and_index
|
||||
from .rag import retrieve, ask_wllm
|
||||
"""Experimental txtai/litellm RAG interface built on the OCR pipeline.
|
||||
|
||||
This package is a prototype. The supported interface is the `ocr-pipeline` CLI
|
||||
(src/ocr_pipeline). Extra dependencies are required: `uv sync --extra kg`.
|
||||
"""
|
||||
|
||||
try:
|
||||
from .ocr import get_screenshots, extract_text
|
||||
from .embeddings import create_and_index
|
||||
from .rag import retrieve, ask_wllm
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"kg_ocr is experimental and needs extra dependencies: "
|
||||
"uv sync --extra kg (or: uv pip install txtai litellm python-dotenv)"
|
||||
) from exc
|
||||
|
||||
__all__ = ["get_screenshots", "extract_text", "create_and_index", "retrieve", "ask_wllm"]
|
||||
|
||||
@@ -16,45 +16,22 @@ dependencies = [
|
||||
"pyyaml>=6.0.1",
|
||||
"structlog>=24.1.0",
|
||||
"rich>=13.7.0",
|
||||
"tqdm>=4.66.0",
|
||||
|
||||
# Image processing
|
||||
"opencv-python-headless>=4.9.0",
|
||||
"opencv-python-headless>=4.9.0,<5.0.0",
|
||||
"pillow>=10.2.0",
|
||||
"numpy>=1.26.0,<2.0; python_version < '3.12'",
|
||||
"numpy>=1.26.0; python_version >= '3.12'",
|
||||
"scikit-image>=0.22.0,<0.23.0; python_version < '3.12'",
|
||||
"scikit-image>=0.22.0; python_version >= '3.12'",
|
||||
|
||||
# OCR engines
|
||||
"paddleocr>=2.7.0,<3.0.0",
|
||||
"paddlepaddle>=2.6.0",
|
||||
"pytesseract>=0.3.10",
|
||||
|
||||
# Table detection
|
||||
"torch>=2.2.0",
|
||||
"torchvision>=0.17.0",
|
||||
"transformers>=4.38.0",
|
||||
"timm>=1.0.0",
|
||||
"layoutparser>=0.3.0",
|
||||
|
||||
# NLP
|
||||
# Text processing and notebook migration
|
||||
"langchain-text-splitters>=0.0.2",
|
||||
"spacy>=3.7.0,<3.8.0; python_version < '3.12'",
|
||||
"spacy>=3.7.0; python_version >= '3.12'",
|
||||
"nbformat>=5.9.0",
|
||||
"legacy-cgi>=2.6.2",
|
||||
|
||||
# Database
|
||||
"sqlite-utils>=3.37.0",
|
||||
|
||||
# File watching
|
||||
"watchdog>=3.0.0",
|
||||
|
||||
# Utilities
|
||||
"python-slugify>=8.0.0",
|
||||
"xxhash>=3.4.0",
|
||||
"python-magic>=0.4.27",
|
||||
"platformdirs>=4.2.0",
|
||||
]
|
||||
|
||||
@@ -67,16 +44,43 @@ dev = [
|
||||
"mypy>=1.8.0",
|
||||
"pre-commit>=3.6.0",
|
||||
]
|
||||
paddle = [
|
||||
"paddleocr>=2.7.0,<3.0.0",
|
||||
"paddlepaddle>=2.6.0",
|
||||
]
|
||||
tables = [
|
||||
"torch>=2.2.0",
|
||||
"torchvision>=0.17.0",
|
||||
"transformers>=4.38.0",
|
||||
"timm>=1.0.0",
|
||||
]
|
||||
figures = [
|
||||
"layoutparser>=0.3.0",
|
||||
"legacy-cgi>=2.6.2; python_version >= '3.13'",
|
||||
]
|
||||
scientific = [
|
||||
"spacy>=3.7.0,<3.8.0; python_version < '3.12'",
|
||||
"spacy>=3.7.0; python_version >= '3.12'",
|
||||
]
|
||||
kg = [
|
||||
"txtai>=7.4.0",
|
||||
"litellm>=1.40.0",
|
||||
"python-dotenv>=1.0.0",
|
||||
]
|
||||
grobid = []
|
||||
full = [
|
||||
"ocr-pipeline[dev]",
|
||||
"ocr-pipeline[paddle,tables,figures,scientific]",
|
||||
]
|
||||
|
||||
[project.entry-points.console_scripts]
|
||||
ocr-pipeline = "ocr_pipeline.cli:app"
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["src"]
|
||||
where = ["src", "."]
|
||||
include = ["ocr_pipeline*", "kg_ocr*"]
|
||||
|
||||
[tool.setuptools.package-data]
|
||||
ocr_pipeline = ["assets/*.png"]
|
||||
|
||||
[tool.uv]
|
||||
dev-dependencies = [
|
||||
|
||||
@@ -4,7 +4,7 @@ Version: 0.1.0
|
||||
Summary: OCR pipeline for life science screenshots - RAG ready
|
||||
Requires-Python: >=3.11
|
||||
Description-Content-Type: text/markdown
|
||||
Requires-Dist: typer[all]>=0.9.0
|
||||
Requires-Dist: typer>=0.9.0
|
||||
Requires-Dist: pydantic>=2.6.0
|
||||
Requires-Dist: pydantic-settings>=2.2.0
|
||||
Requires-Dist: pyyaml>=6.0.1
|
||||
@@ -13,20 +13,11 @@ Requires-Dist: rich>=13.7.0
|
||||
Requires-Dist: tqdm>=4.66.0
|
||||
Requires-Dist: opencv-python-headless>=4.9.0
|
||||
Requires-Dist: pillow>=10.2.0
|
||||
Requires-Dist: numpy>=1.26.0
|
||||
Requires-Dist: scikit-image>=0.22.0
|
||||
Requires-Dist: paddleocr<3.0.0,>=2.7.0
|
||||
Requires-Dist: paddlepaddle>=2.6.0
|
||||
Requires-Dist: numpy<2.0,>=1.26.0; python_version < "3.12"
|
||||
Requires-Dist: numpy>=1.26.0; python_version >= "3.12"
|
||||
Requires-Dist: pytesseract>=0.3.10
|
||||
Requires-Dist: torch>=2.2.0
|
||||
Requires-Dist: torchvision>=0.17.0
|
||||
Requires-Dist: transformers>=4.38.0
|
||||
Requires-Dist: timm>=1.0.0
|
||||
Requires-Dist: layoutparser>=0.3.0
|
||||
Requires-Dist: langchain-text-splitters>=0.0.2
|
||||
Requires-Dist: spacy>=3.7.0
|
||||
Requires-Dist: nbformat>=5.9.0
|
||||
Requires-Dist: legacy-cgi>=2.6.2
|
||||
Requires-Dist: sqlite-utils>=3.37.0
|
||||
Requires-Dist: watchdog>=3.0.0
|
||||
Requires-Dist: python-slugify>=8.0.0
|
||||
@@ -41,32 +32,92 @@ Requires-Dist: ruff>=0.2.0; extra == "dev"
|
||||
Requires-Dist: mypy>=1.8.0; extra == "dev"
|
||||
Requires-Dist: pre-commit>=3.6.0; extra == "dev"
|
||||
Provides-Extra: grobid
|
||||
Provides-Extra: paddle
|
||||
Requires-Dist: paddleocr<3.0.0,>=2.7.0; extra == "paddle"
|
||||
Requires-Dist: paddlepaddle>=2.6.0; extra == "paddle"
|
||||
Provides-Extra: scientific
|
||||
Requires-Dist: spacy<3.8.0,>=3.7.0; python_version < "3.12" and extra == "scientific"
|
||||
Requires-Dist: spacy>=3.7.0; python_version >= "3.12" and extra == "scientific"
|
||||
Requires-Dist: scispacy==0.5.4; python_version < "3.12" and extra == "scientific"
|
||||
Provides-Extra: tables
|
||||
Requires-Dist: torch>=2.2.0; extra == "tables"
|
||||
Requires-Dist: torchvision>=0.17.0; extra == "tables"
|
||||
Requires-Dist: transformers>=4.38.0; extra == "tables"
|
||||
Requires-Dist: timm>=1.0.0; extra == "tables"
|
||||
Provides-Extra: figures
|
||||
Requires-Dist: layoutparser>=0.3.0; extra == "figures"
|
||||
Provides-Extra: full
|
||||
Requires-Dist: ocr-pipeline[dev]; extra == "full"
|
||||
Requires-Dist: ocr-pipeline[dev,figures,paddle,scientific,tables]; extra == "full"
|
||||
|
||||
# OCR Pipeline for Life Science Screenshots
|
||||
|
||||
Turns scientific screenshots into RAG-ready Markdown, with metadata attached.
|
||||
|
||||
- OCR via PaddleOCR, falls back to Tesseract
|
||||
- Reliable Tesseract OCR by default, with optional PaddleOCR
|
||||
- Preprocessing: deskew, denoise, CLAHE, line removal
|
||||
- Figure/table detection
|
||||
- Entity extraction with scispaCy
|
||||
- Citation matching
|
||||
- Chunking for embedding
|
||||
- Optional figure/table detection
|
||||
- Optional scientific entity extraction with scispaCy
|
||||
- Citation matching and embedding-ready chunking
|
||||
|
||||
## Quick start
|
||||
|
||||
```bash
|
||||
uv sync # or: pip install -e .
|
||||
pre-commit install
|
||||
# The default basic profile is lightweight: Tesseract, cleanup, citations,
|
||||
# chunking, Markdown, and no model downloads.
|
||||
uv sync
|
||||
uv run ocr-pipeline doctor
|
||||
|
||||
ocr-pipeline run # process screenshots
|
||||
ocr-pipeline watch # watch a folder
|
||||
ocr-pipeline reprocess ocr_sc.ipynb # pull screenshots out of a notebook
|
||||
ocr-pipeline status
|
||||
# Smoke-test a small directory.
|
||||
uv run ocr-pipeline run \
|
||||
--input-dir ~/Pictures/test_screenshots \
|
||||
--output-dir /tmp/ocr-output \
|
||||
--workers 1 \
|
||||
--exclude "*.photoslibrary/*"
|
||||
```
|
||||
|
||||
Use `ocr-pipeline doctor` before a first run or after changing environments. It checks the active profile without loading ML models or downloading model weights.
|
||||
|
||||
Use `--force` to reprocess files already recorded in the processing database. The
|
||||
`--exclude` option is repeatable and prevents recursive scans from entering bundles
|
||||
such as macOS `.photoslibrary` directories.
|
||||
|
||||
### Capability profiles
|
||||
|
||||
| Profile | What it enables | Intended environment |
|
||||
| --- | --- | --- |
|
||||
| `basic` (default) | Tesseract OCR, cleanup, citations, chunks, Markdown | Python 3.11+ |
|
||||
| `scientific` | Basic + `en_core_sci_lg` entity extraction | **Python 3.11** |
|
||||
| `full` | Scientific + PaddleOCR + figure/table detection | Platform-specific ML environment |
|
||||
|
||||
Select a profile per run with `--profile`, or set `profile: basic`, `scientific`, or
|
||||
`full` at the top of `config.yaml`. The profile is an explicit promise: unavailable
|
||||
optional features are reported by `doctor`; they are not silently initialized on a
|
||||
basic run.
|
||||
|
||||
```bash
|
||||
uv run ocr-pipeline doctor --profile scientific
|
||||
uv run ocr-pipeline setup scientific # prints the safe setup commands
|
||||
uv run ocr-pipeline run --profile scientific --input-dir ~/Pictures/test_screenshots
|
||||
```
|
||||
|
||||
#### Scientific NER environment
|
||||
|
||||
`en_core_sci_lg` is distributed separately from scispaCy. Create a dedicated Python
|
||||
3.11 environment—do not mix its compiled spaCy stack with a Python 3.13 environment:
|
||||
|
||||
```bash
|
||||
uv venv --python 3.11 .venv-scientific
|
||||
uv pip install --python .venv-scientific/bin/python -e '.[scientific]'
|
||||
uv pip install --python .venv-scientific/bin/python \
|
||||
https://s3-us-west-2.amazonaws.com/ai2-s2-scispacy/releases/v0.5.4/en_core_sci_lg-0.5.4.tar.gz
|
||||
|
||||
.venv-scientific/bin/ocr-pipeline doctor --profile scientific
|
||||
.venv-scientific/bin/ocr-pipeline run --profile scientific --input-dir ~/Pictures/test_screenshots
|
||||
```
|
||||
|
||||
Use the explicit `.venv-scientific/bin/ocr-pipeline` path for that environment;
|
||||
`uv run` always selects the project's default `.venv`.
|
||||
|
||||
## Docker
|
||||
|
||||
```bash
|
||||
@@ -86,9 +137,10 @@ input:
|
||||
paths: ["~/Pictures", "/mnt/storage3/aman/screenshots"]
|
||||
patterns: ["SCR-*.png", "*.jpg", "*.jpeg", "*.tiff"]
|
||||
recursive: true
|
||||
exclude_patterns: ["*.photoslibrary/*"]
|
||||
|
||||
ocr:
|
||||
engine: "paddleocr" # paddleocr | tesseract | auto
|
||||
engine: "tesseract" # paddleocr | tesseract | auto
|
||||
languages: ["en", "latin"]
|
||||
use_gpu: false
|
||||
preprocess:
|
||||
@@ -105,14 +157,14 @@ processing:
|
||||
|
||||
detectors:
|
||||
figures:
|
||||
enabled: true
|
||||
enabled: false
|
||||
confidence_threshold: 0.7
|
||||
tables:
|
||||
enabled: true
|
||||
enabled: false
|
||||
confidence_threshold: 0.7
|
||||
|
||||
entities:
|
||||
enabled: true
|
||||
enabled: false
|
||||
model: "en_core_sci_lg"
|
||||
|
||||
citations:
|
||||
@@ -126,6 +178,8 @@ chunking:
|
||||
output:
|
||||
base_directory: "./data/ocr_output"
|
||||
organize_by: "date_run" # date_run | source_dir | flat
|
||||
write_consolidated: true
|
||||
consolidated_filename: "all_ocr.md"
|
||||
frontmatter:
|
||||
- source_path
|
||||
- source_hash
|
||||
@@ -134,6 +188,7 @@ output:
|
||||
- ocr_confidence_mean
|
||||
- language
|
||||
- detected_entities
|
||||
- entity_extraction_backend
|
||||
- has_figures
|
||||
- has_tables
|
||||
- citations_found
|
||||
@@ -200,15 +255,16 @@ Lane 4: BRCA1 KO + pBRCA1-C61G mutant
|
||||
Anti-BRCA1 (1:1000), Anti-β-actin (1:5000)
|
||||
```
|
||||
|
||||
Files land in `data/ocr_output/<date>/run_NNN/`.
|
||||
Files land in `data/ocr_output/<date>/run_NNN/`. Individual chunk files are retained for RAG indexing. Set `output.write_consolidated: true` to also write `all_ocr.md` containing one frontmatter block and all source text grouped by image.
|
||||
|
||||
## Life science specifics
|
||||
|
||||
Gene/protein names go through scispaCy's `en_core_sci_lg`. Chemical formulas and units (µM, ng/mL, kb/Mb/Gb, °C, ×g) get normalized, scientific notation gets cleaned up (`1.5×10⁻³` → `1.5×10^-3`), and gel/blot figures get their captions pulled out separately. Citations are matched by regex for DOI, PMID, arXiv, PMC, and ISBN.
|
||||
|
||||
## Models
|
||||
## Optional models
|
||||
|
||||
First run downloads what it needs, cached in `~/.cache/ocr_pipeline/`:
|
||||
The basic profile downloads no ML models. Advanced profiles download models only when
|
||||
the corresponding capability is enabled and used, cached in `~/.cache/ocr_pipeline/`:
|
||||
|
||||
- PaddleOCR models (~200MB)
|
||||
- scispaCy `en_core_sci_lg` (~800MB)
|
||||
@@ -224,7 +280,7 @@ The core pipeline can run with Tesseract alone. PaddleOCR, scientific NER, figur
|
||||
- **Figures:** `layoutparser`'s `Detectron2LayoutModel` requires a Detectron2 build matching your Torch/Python platform. It is intentionally not forced as a universal dependency because no single wheel supports every platform.
|
||||
- **Tables:** the Table Transformer model is downloaded by Transformers on first use; ensure the selected model's optional dependencies (including `timm`, when required by that model revision) are installed in the runtime image.
|
||||
|
||||
Use `ocr-pipeline config` to verify effective settings. A missing optional model is logged and produces empty results for that detector instead of failing a complete batch.
|
||||
Use `ocr-pipeline doctor --profile <profile>` to verify dependencies before processing. Use `ocr-pipeline config` to inspect effective settings. A configured but unavailable optional model is reported as missing by `doctor`; basic runs never attempt to initialize it.
|
||||
|
||||
## TODO
|
||||
|
||||
|
||||
@@ -4,6 +4,7 @@ setup.py
|
||||
src/ocr_pipeline/__init__.py
|
||||
src/ocr_pipeline/cli.py
|
||||
src/ocr_pipeline/config.py
|
||||
src/ocr_pipeline/doctor.py
|
||||
src/ocr_pipeline/pipeline.py
|
||||
src/ocr_pipeline/watch.py
|
||||
src/ocr_pipeline.egg-info/PKG-INFO
|
||||
@@ -38,4 +39,5 @@ tests/test_graph.py
|
||||
tests/test_indexer.py
|
||||
tests/test_ocr.py
|
||||
tests/test_output.py
|
||||
tests/test_postprocess.py
|
||||
tests/test_postprocess.py
|
||||
tests/test_profiles.py
|
||||
@@ -1,4 +1,4 @@
|
||||
typer[all]>=0.9.0
|
||||
typer>=0.9.0
|
||||
pydantic>=2.6.0
|
||||
pydantic-settings>=2.2.0
|
||||
pyyaml>=6.0.1
|
||||
@@ -7,20 +7,9 @@ rich>=13.7.0
|
||||
tqdm>=4.66.0
|
||||
opencv-python-headless>=4.9.0
|
||||
pillow>=10.2.0
|
||||
numpy>=1.26.0
|
||||
scikit-image>=0.22.0
|
||||
paddleocr<3.0.0,>=2.7.0
|
||||
paddlepaddle>=2.6.0
|
||||
pytesseract>=0.3.10
|
||||
torch>=2.2.0
|
||||
torchvision>=0.17.0
|
||||
transformers>=4.38.0
|
||||
timm>=1.0.0
|
||||
layoutparser>=0.3.0
|
||||
langchain-text-splitters>=0.0.2
|
||||
spacy>=3.7.0
|
||||
nbformat>=5.9.0
|
||||
legacy-cgi>=2.6.2
|
||||
sqlite-utils>=3.37.0
|
||||
watchdog>=3.0.0
|
||||
python-slugify>=8.0.0
|
||||
@@ -28,6 +17,12 @@ xxhash>=3.4.0
|
||||
python-magic>=0.4.27
|
||||
platformdirs>=4.2.0
|
||||
|
||||
[:python_version < "3.12"]
|
||||
numpy<2.0,>=1.26.0
|
||||
|
||||
[:python_version >= "3.12"]
|
||||
numpy>=1.26.0
|
||||
|
||||
[dev]
|
||||
pytest>=8.0.0
|
||||
pytest-cov>=4.1.0
|
||||
@@ -36,7 +31,29 @@ ruff>=0.2.0
|
||||
mypy>=1.8.0
|
||||
pre-commit>=3.6.0
|
||||
|
||||
[figures]
|
||||
layoutparser>=0.3.0
|
||||
|
||||
[full]
|
||||
ocr-pipeline[dev]
|
||||
ocr-pipeline[dev,figures,paddle,scientific,tables]
|
||||
|
||||
[grobid]
|
||||
|
||||
[paddle]
|
||||
paddleocr<3.0.0,>=2.7.0
|
||||
paddlepaddle>=2.6.0
|
||||
|
||||
[scientific]
|
||||
|
||||
[scientific:python_version < "3.12"]
|
||||
spacy<3.8.0,>=3.7.0
|
||||
scispacy==0.5.4
|
||||
|
||||
[scientific:python_version >= "3.12"]
|
||||
spacy>=3.7.0
|
||||
|
||||
[tables]
|
||||
torch>=2.2.0
|
||||
torchvision>=0.17.0
|
||||
transformers>=4.38.0
|
||||
timm>=1.0.0
|
||||
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
0
src/ocr_pipeline/assets/__init__.py
Normal file
0
src/ocr_pipeline/assets/__init__.py
Normal file
BIN
src/ocr_pipeline/assets/demo_screenshot.png
Normal file
BIN
src/ocr_pipeline/assets/demo_screenshot.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 60 KiB |
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -13,4 +13,10 @@ from .regex_extractor import (
|
||||
format_citations_for_markdown,
|
||||
)
|
||||
|
||||
__all__ = ["Citation", "GrobidClient", "RegexCitationExtractor", "extract_citations", "format_citations_for_markdown"]
|
||||
__all__ = [
|
||||
"Citation",
|
||||
"GrobidClient",
|
||||
"RegexCitationExtractor",
|
||||
"extract_citations",
|
||||
"format_citations_for_markdown",
|
||||
]
|
||||
|
||||
@@ -2,13 +2,22 @@ from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Literal, cast
|
||||
|
||||
import typer
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
|
||||
from ocr_pipeline.config import settings
|
||||
from ocr_pipeline.config import find_config_file, settings
|
||||
from ocr_pipeline.demo import run_demo
|
||||
from ocr_pipeline.onboarding import (
|
||||
default_config_location,
|
||||
tesseract_install_hint,
|
||||
tesseract_path,
|
||||
write_default_config,
|
||||
)
|
||||
from ocr_pipeline.pipeline import OCRPipeline
|
||||
from ocr_pipeline.setup_steps import run_steps, steps_for
|
||||
from ocr_pipeline.utils.logging import get_logger, setup_logging
|
||||
from ocr_pipeline.utils.migrate import migrate_notebook
|
||||
|
||||
@@ -21,6 +30,18 @@ console = Console()
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def _ensure_configured(input_dir: Path | None) -> None:
|
||||
"""Stop early with guidance instead of scanning a default directory."""
|
||||
if input_dir is not None or find_config_file() is not None:
|
||||
return
|
||||
console.print("[yellow]No config.yaml found and no --input-dir given.[/yellow]")
|
||||
console.print("Get started:")
|
||||
console.print(" [bold]ocr-pipeline init[/bold] create a config for this machine")
|
||||
console.print(" [bold]ocr-pipeline demo[/bold] try the pipeline on a sample screenshot")
|
||||
console.print(" [bold]ocr-pipeline run -i DIR[/bold] point at a directory directly")
|
||||
raise typer.Exit(code=2)
|
||||
|
||||
|
||||
@app.callback()
|
||||
def callback(
|
||||
config: Path | None = typer.Option(None, "--config", "-c", help="Config file path"),
|
||||
@@ -31,23 +52,92 @@ def callback(
|
||||
setup_logging(log_level, log_file)
|
||||
|
||||
if config:
|
||||
if not config.exists():
|
||||
console.print(f"[red]Config file not found:[/red] {config}")
|
||||
raise typer.Exit(code=2)
|
||||
settings.update_from_yaml(config)
|
||||
|
||||
|
||||
@app.command()
|
||||
def init(
|
||||
local: bool = typer.Option(
|
||||
False, "--local", help="Write ./config.yaml instead of the per-user config directory"
|
||||
),
|
||||
input_dir: Path | None = typer.Option(
|
||||
None, "--input-dir", "-i", help="Screenshots directory to put in the config"
|
||||
),
|
||||
force: bool = typer.Option(False, "--force", help="Overwrite an existing config file"),
|
||||
):
|
||||
"""Create a minimal config for this machine."""
|
||||
target_dir = default_config_location(local)
|
||||
config_path = target_dir / "config.yaml"
|
||||
if config_path.exists() and not force:
|
||||
console.print(f"[yellow]Config already exists:[/yellow] {config_path}")
|
||||
console.print("Use --force to overwrite it.")
|
||||
raise typer.Exit(code=1)
|
||||
|
||||
written = write_default_config(target_dir, input_dir)
|
||||
console.print(f"[green]Config written:[/green] {written}")
|
||||
|
||||
tess = tesseract_path()
|
||||
if tess:
|
||||
console.print(f"[green]Tesseract found:[/green] {tess}")
|
||||
else:
|
||||
console.print("[yellow]Tesseract not found on PATH.[/yellow]")
|
||||
console.print(f" Install it with: [bold]{tesseract_install_hint()}[/bold]")
|
||||
console.print(" or run: [bold]ocr-pipeline setup[/bold]")
|
||||
|
||||
console.print("\nNext steps:")
|
||||
console.print(
|
||||
" [bold]ocr-pipeline demo[/bold] verify everything works on a sample screenshot"
|
||||
)
|
||||
console.print(" [bold]ocr-pipeline run[/bold] process your screenshots")
|
||||
|
||||
|
||||
@app.command()
|
||||
def demo():
|
||||
"""Run the pipeline on a bundled sample screenshot to verify the install."""
|
||||
console.print("[green]Running demo on the bundled sample screenshot...[/green]")
|
||||
result = run_demo()
|
||||
|
||||
table = Table(title="Demo Results")
|
||||
table.add_column("Metric", style="cyan")
|
||||
table.add_column("Value", style="green")
|
||||
table.add_row("Successful", str(result.successful))
|
||||
table.add_row("Failed", str(result.failed))
|
||||
table.add_row("Chunks Created", str(result.chunks_created))
|
||||
console.print(table)
|
||||
|
||||
if result.failed > 0:
|
||||
for err in result.errors:
|
||||
console.print(f" [red]{err['path']}: {err['error']}[/red]")
|
||||
console.print("Something is wrong with the install; run [bold]ocr-pipeline setup[/bold].")
|
||||
raise typer.Exit(code=1)
|
||||
|
||||
markdown_files = [path for path in result.output_files if path.suffix == ".md"]
|
||||
if markdown_files:
|
||||
first = markdown_files[0]
|
||||
console.print(f"\n[green]Demo output:[/green] {first}")
|
||||
excerpt = first.read_text(encoding="utf-8")[:1200]
|
||||
console.print(f"\n[dim]{excerpt}[/dim]")
|
||||
console.print("\n[green]Install verified.[/green] Run [bold]ocr-pipeline init[/bold] next.")
|
||||
|
||||
|
||||
@app.command()
|
||||
def run(
|
||||
input_dir: Path | None = typer.Option(None, "--input-dir", "-i", help="Input directory"),
|
||||
output_dir: Path | None = typer.Option(None, "--output-dir", "-o", help="Output directory"),
|
||||
workers: int | None = typer.Option(
|
||||
None, "--workers", "-w", help="Number of worker processes"
|
||||
),
|
||||
workers: int | None = typer.Option(None, "--workers", "-w", help="Number of worker processes"),
|
||||
engine: str | None = typer.Option(
|
||||
None, "--engine", "-e", help="OCR engine (paddleocr/tesseract/auto)"
|
||||
),
|
||||
exclude: list[str] = typer.Option([], "--exclude", help="Glob pattern to exclude; repeatable"),
|
||||
force: bool = typer.Option(False, "--force", help="Reprocess files already recorded as processed"),
|
||||
force: bool = typer.Option(
|
||||
False, "--force", help="Reprocess files already recorded as processed"
|
||||
),
|
||||
):
|
||||
"""Run OCR pipeline on all screenshots"""
|
||||
_ensure_configured(input_dir)
|
||||
if input_dir:
|
||||
settings.input.paths = [str(input_dir)]
|
||||
if output_dir:
|
||||
@@ -55,7 +145,12 @@ def run(
|
||||
if workers:
|
||||
settings.processing.workers = workers
|
||||
if engine:
|
||||
settings.ocr.engine = engine
|
||||
if engine not in ("paddleocr", "tesseract", "auto"):
|
||||
console.print(
|
||||
f"[red]Unknown engine:[/red] {engine} (choose paddleocr, tesseract, or auto)"
|
||||
)
|
||||
raise typer.Exit(code=2)
|
||||
settings.ocr.engine = cast(Literal["paddleocr", "tesseract", "auto"], engine)
|
||||
if exclude:
|
||||
settings.input.exclude_patterns = exclude
|
||||
if force:
|
||||
@@ -70,6 +165,7 @@ def run(
|
||||
|
||||
table.add_row("Total Images", str(result.total_images))
|
||||
table.add_row("Successful", str(result.successful))
|
||||
table.add_row("Skipped (already processed)", str(result.skipped))
|
||||
table.add_row("Failed", str(result.failed))
|
||||
table.add_row("Chunks Created", str(result.chunks_created))
|
||||
table.add_row("Output Files", str(len(result.output_files)))
|
||||
@@ -88,6 +184,7 @@ def run(
|
||||
@app.command()
|
||||
def watch():
|
||||
"""Watch for new screenshots and process automatically"""
|
||||
_ensure_configured(None)
|
||||
pipeline = OCRPipeline()
|
||||
console.print("[green]Starting watch mode...[/green]")
|
||||
console.print("Press Ctrl+C to stop")
|
||||
@@ -154,9 +251,35 @@ def stats(days: int = typer.Option(30, "--days", "-d", help="Number of days to a
|
||||
console.print(table)
|
||||
|
||||
|
||||
@app.command()
|
||||
def setup(
|
||||
target: str = typer.Argument("basic", help="What to set up: basic or full"),
|
||||
yes: bool = typer.Option(False, "--yes", "-y", help="Run all steps without asking"),
|
||||
dry_run: bool = typer.Option(False, "--dry-run", help="Print the steps without running them"),
|
||||
):
|
||||
"""Install system/optional dependencies, executing each step with confirmation."""
|
||||
try:
|
||||
steps = steps_for(target)
|
||||
except ValueError as exc:
|
||||
console.print(f"[red]{exc}[/red]")
|
||||
raise typer.Exit(code=2) from exc
|
||||
if dry_run:
|
||||
console.print(f"[cyan]Setup steps for {target} (dry run):[/cyan]")
|
||||
returncode = run_steps(steps, assume_yes=yes, dry_run=dry_run, console=console)
|
||||
if returncode != 0:
|
||||
raise typer.Exit(code=returncode)
|
||||
if not dry_run:
|
||||
console.print("[green]Setup complete.[/green]")
|
||||
|
||||
|
||||
@app.command()
|
||||
def config():
|
||||
"""Show current configuration"""
|
||||
source = find_config_file()
|
||||
if source:
|
||||
console.print(f"[dim]Loaded from: {source}[/dim]")
|
||||
else:
|
||||
console.print("[dim]No config.yaml found; using built-in defaults.[/dim]")
|
||||
console.print("[cyan]Current Configuration:[/cyan]")
|
||||
console.print(settings.model_dump_json(indent=2))
|
||||
|
||||
|
||||
@@ -3,7 +3,7 @@ from __future__ import annotations
|
||||
import os
|
||||
import platform
|
||||
from pathlib import Path
|
||||
from typing import Any, ClassVar, Literal
|
||||
from typing import Any, ClassVar, Literal, cast
|
||||
|
||||
import platformdirs
|
||||
import yaml
|
||||
@@ -12,6 +12,8 @@ from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
APP_NAME = "ocr-pipeline"
|
||||
APP_AUTHOR = "aman"
|
||||
CONFIG_ENV_VAR = "OCR_PIPELINE_CONFIG"
|
||||
DEFAULT_CONFIG_FILENAME = "config.yaml"
|
||||
|
||||
|
||||
def _get_platform_ignore_patterns() -> list[str]:
|
||||
@@ -60,7 +62,7 @@ class PreprocessConfig(BaseSettings):
|
||||
|
||||
|
||||
class OCRConfig(BaseSettings):
|
||||
engine: Literal["paddleocr", "tesseract", "auto"] = "paddleocr"
|
||||
engine: Literal["paddleocr", "tesseract", "auto"] = "tesseract"
|
||||
languages: list[str] = Field(default_factory=lambda: ["en", "latin"])
|
||||
use_gpu: bool = False
|
||||
use_angle_cls: bool = True
|
||||
@@ -81,7 +83,7 @@ class ProcessingConfig(BaseSettings):
|
||||
|
||||
|
||||
class DetectorConfig(BaseSettings):
|
||||
enabled: bool = True
|
||||
enabled: bool = False
|
||||
model: str = ""
|
||||
confidence_threshold: float = Field(default=0.7, ge=0.0, le=1.0)
|
||||
|
||||
@@ -108,7 +110,7 @@ class DetectorsConfig(BaseSettings):
|
||||
|
||||
|
||||
class EntitiesConfig(BaseSettings):
|
||||
enabled: bool = True
|
||||
enabled: bool = False
|
||||
model: str = "en_core_sci_lg"
|
||||
types: list[str] = Field(
|
||||
default_factory=lambda: [
|
||||
@@ -176,7 +178,9 @@ class WatchConfig(BaseSettings):
|
||||
debounce_seconds: int = Field(default=5, ge=0)
|
||||
ignore_patterns: list[str] = Field(default_factory=_get_platform_ignore_patterns)
|
||||
db_path: str = Field(
|
||||
default_factory=lambda: str(Path(platformdirs.user_data_dir(APP_NAME, APP_AUTHOR)) / "processed_files.db")
|
||||
default_factory=lambda: str(
|
||||
Path(platformdirs.user_data_dir(APP_NAME, APP_AUTHOR)) / "processed_files.db"
|
||||
)
|
||||
)
|
||||
poll_interval: float = Field(default=1.0, gt=0)
|
||||
|
||||
@@ -189,15 +193,45 @@ class _BaseSettings(BaseSettings):
|
||||
)
|
||||
|
||||
|
||||
def _load_yaml_config(path: str = "config.yaml") -> dict[str, Any]:
|
||||
"""Load YAML config file if it exists."""
|
||||
config_path = Path(path)
|
||||
if not config_path.is_absolute():
|
||||
config_path = Path.cwd() / config_path
|
||||
if config_path.exists():
|
||||
with open(config_path) as f:
|
||||
return yaml.safe_load(f) or {}
|
||||
return {}
|
||||
def user_config_dir() -> Path:
|
||||
"""Per-user config directory for this machine."""
|
||||
return Path(platformdirs.user_config_dir(APP_NAME, APP_AUTHOR))
|
||||
|
||||
|
||||
def _candidate_config_paths() -> list[Path]:
|
||||
"""Config file locations in priority order."""
|
||||
candidates: list[Path] = []
|
||||
env_path = os.environ.get(CONFIG_ENV_VAR)
|
||||
if env_path:
|
||||
candidates.append(Path(env_path).expanduser())
|
||||
candidates.append(Path.cwd() / DEFAULT_CONFIG_FILENAME)
|
||||
candidates.append(user_config_dir() / DEFAULT_CONFIG_FILENAME)
|
||||
return candidates
|
||||
|
||||
|
||||
def find_config_file() -> Path | None:
|
||||
"""First config file that exists, or None when nothing is configured."""
|
||||
for candidate in _candidate_config_paths():
|
||||
if candidate.is_file():
|
||||
return candidate
|
||||
return None
|
||||
|
||||
|
||||
def _load_yaml_config(path: str | Path | None = None) -> dict[str, Any]:
|
||||
"""Load YAML config from an explicit path or the standard search paths."""
|
||||
if path is not None:
|
||||
config_path = Path(path)
|
||||
if not config_path.is_absolute():
|
||||
config_path = Path.cwd() / config_path
|
||||
if config_path.exists():
|
||||
with open(config_path) as f:
|
||||
return yaml.safe_load(f) or {}
|
||||
return {}
|
||||
found = find_config_file()
|
||||
if found is None:
|
||||
return {}
|
||||
with open(found) as f:
|
||||
return yaml.safe_load(f) or {}
|
||||
|
||||
|
||||
def _build_settings_from_yaml(yaml_dict: dict[str, Any]) -> dict[str, Any]:
|
||||
@@ -253,7 +287,7 @@ class Settings(_BaseSettings):
|
||||
yaml_models = _build_settings_from_yaml(Settings._yaml_config)
|
||||
super().__init__(**yaml_models, **kwargs)
|
||||
|
||||
def update_from_yaml(self, path: str) -> None:
|
||||
def update_from_yaml(self, path: str | Path) -> None:
|
||||
"""Reload settings from a YAML file."""
|
||||
Settings._yaml_config = _load_yaml_config(path)
|
||||
# Re-initialize with new config
|
||||
@@ -271,4 +305,34 @@ class Settings(_BaseSettings):
|
||||
return Path(self.watch.db_path).expanduser().resolve()
|
||||
|
||||
|
||||
settings = Settings()
|
||||
class _LazySettings:
|
||||
"""Import-safe handle for the global Settings.
|
||||
|
||||
Instantiating Settings at import time read config.yaml from whatever
|
||||
directory the process started in, so `import ocr_pipeline` touched the
|
||||
filesystem and behaved differently depending on the caller's cwd. This
|
||||
proxy defers all of that until the settings are actually used.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._instance: Settings | None = None
|
||||
|
||||
def _get_instance(self) -> Settings:
|
||||
if self._instance is None:
|
||||
self._instance = Settings()
|
||||
return self._instance
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
return getattr(self._get_instance(), name)
|
||||
|
||||
def __setattr__(self, name: str, value: Any) -> None:
|
||||
if name == "_instance":
|
||||
object.__setattr__(self, name, value)
|
||||
else:
|
||||
setattr(self._get_instance(), name, value)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return repr(self._get_instance())
|
||||
|
||||
|
||||
settings: Settings = cast(Settings, _LazySettings())
|
||||
|
||||
45
src/ocr_pipeline/demo.py
Normal file
45
src/ocr_pipeline/demo.py
Normal file
@@ -0,0 +1,45 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import shutil
|
||||
import tempfile
|
||||
from importlib.resources import as_file, files
|
||||
from pathlib import Path
|
||||
|
||||
from ocr_pipeline.config import settings
|
||||
from ocr_pipeline.pipeline import OCRPipeline, PipelineResult
|
||||
from ocr_pipeline.utils.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
DEMO_IMAGE_NAME = "demo_screenshot.png"
|
||||
|
||||
|
||||
def demo_image_path() -> Path:
|
||||
"""Copy the bundled sample screenshot into a fresh temp directory."""
|
||||
resource = files("ocr_pipeline.assets") / DEMO_IMAGE_NAME
|
||||
target_dir = Path(tempfile.mkdtemp(prefix="ocr-pipeline-demo-"))
|
||||
target = target_dir / DEMO_IMAGE_NAME
|
||||
with as_file(resource) as source:
|
||||
shutil.copy(source, target)
|
||||
return target
|
||||
|
||||
|
||||
def _apply_demo_settings(work_dir: Path) -> None:
|
||||
"""Hermetic demo: guaranteed-safe engine and features, isolated output and DB."""
|
||||
settings.ocr.engine = "tesseract"
|
||||
settings.entities.enabled = False
|
||||
settings.detectors.figures.enabled = False
|
||||
settings.detectors.tables.enabled = False
|
||||
settings.detectors.captions.enabled = False
|
||||
settings.processing.skip_existing = False
|
||||
settings.output.base_directory = str(work_dir / "output")
|
||||
settings.output.write_consolidated = False
|
||||
settings.watch.db_path = str(work_dir / "demo_processed.db")
|
||||
|
||||
|
||||
def run_demo() -> PipelineResult:
|
||||
"""Run the full pipeline on the bundled sample screenshot."""
|
||||
image = demo_image_path()
|
||||
_apply_demo_settings(image.parent)
|
||||
logger.info("demo_started", image=str(image))
|
||||
return OCRPipeline().process_all([image])
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -33,9 +33,13 @@ class FigureDetector:
|
||||
try:
|
||||
import layoutparser as lp
|
||||
|
||||
model_class = getattr(lp, "Detectron2LayoutModel", None) or getattr(lp, "AutoLayoutModel", None)
|
||||
model_class = getattr(lp, "Detectron2LayoutModel", None) or getattr(
|
||||
lp, "AutoLayoutModel", None
|
||||
)
|
||||
if model_class is None:
|
||||
raise RuntimeError("layoutparser has no compatible layout model backend; install Detectron2")
|
||||
raise RuntimeError(
|
||||
"layoutparser has no compatible layout model backend; install Detectron2"
|
||||
)
|
||||
self._predictor = model_class(
|
||||
config_path=settings.detectors.figures.model,
|
||||
label_map={0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"},
|
||||
@@ -104,9 +108,13 @@ class CaptionExtractor:
|
||||
try:
|
||||
import layoutparser as lp
|
||||
|
||||
model_class = getattr(lp, "Detectron2LayoutModel", None) or getattr(lp, "AutoLayoutModel", None)
|
||||
model_class = getattr(lp, "Detectron2LayoutModel", None) or getattr(
|
||||
lp, "AutoLayoutModel", None
|
||||
)
|
||||
if model_class is None:
|
||||
raise RuntimeError("layoutparser has no compatible layout model backend; install Detectron2")
|
||||
raise RuntimeError(
|
||||
"layoutparser has no compatible layout model backend; install Detectron2"
|
||||
)
|
||||
self._predictor = model_class(
|
||||
config_path=settings.detectors.captions.model,
|
||||
label_map={0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"},
|
||||
|
||||
@@ -80,7 +80,12 @@ class TableDetector:
|
||||
):
|
||||
resized_box = [float(x) for x in box]
|
||||
scale_x, scale_y = original_width / 800, original_height / 800
|
||||
box = [resized_box[0] * scale_x, resized_box[1] * scale_y, resized_box[2] * scale_x, resized_box[3] * scale_y]
|
||||
box = [
|
||||
resized_box[0] * scale_x,
|
||||
resized_box[1] * scale_y,
|
||||
resized_box[2] * scale_x,
|
||||
resized_box[3] * scale_y,
|
||||
]
|
||||
x1, y1, x2, y2 = map(int, box)
|
||||
x1, x2 = max(0, x1), min(original_width, x2)
|
||||
y1, y2 = max(0, y1), min(original_height, y2)
|
||||
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -68,7 +68,7 @@ class OCRResult:
|
||||
|
||||
class PaddleOCREngine:
|
||||
def __init__(self):
|
||||
self._ocr = None
|
||||
self._ocr: Any = None
|
||||
self._initialized = False
|
||||
|
||||
def _init(self):
|
||||
@@ -104,6 +104,8 @@ class PaddleOCREngine:
|
||||
|
||||
def process(self, image: np.ndarray) -> OCRResult:
|
||||
self._init()
|
||||
if self._ocr is None:
|
||||
raise RuntimeError("PaddleOCR failed to initialize")
|
||||
start = time.time()
|
||||
|
||||
result = self._ocr.ocr(image, cls=True)
|
||||
@@ -211,24 +213,25 @@ class OCREngine:
|
||||
self.tesseract = TesseractEngine()
|
||||
self.engine_preference = settings.ocr.engine
|
||||
|
||||
def process(self, image: np.ndarray) -> OCRResult:
|
||||
"""Process a pre-loaded image array."""
|
||||
if self.engine_preference == "paddleocr":
|
||||
try:
|
||||
return self.paddle.process(image)
|
||||
except Exception as e:
|
||||
logger.warning("paddleocr_failed_fallback", error=str(e))
|
||||
return self.tesseract.process(image)
|
||||
def process(self, original: np.ndarray, preprocessed: np.ndarray | None = None) -> OCRResult:
|
||||
"""Process an image, giving each engine the variant it handles best.
|
||||
|
||||
elif self.engine_preference == "tesseract":
|
||||
return self.tesseract.process(image)
|
||||
PaddleOCR's detection handles raw screenshots (including dark mode) far
|
||||
better than the binarized preprocessing derivative, so it receives the
|
||||
original image. Tesseract benefits from binarization and receives the
|
||||
preprocessed derivative; it is also the fallback when Paddle fails.
|
||||
"""
|
||||
if preprocessed is None:
|
||||
preprocessed = original
|
||||
|
||||
else:
|
||||
try:
|
||||
return self.paddle.process(image)
|
||||
except Exception as e:
|
||||
logger.warning("paddleocr_failed_fallback", error=str(e))
|
||||
return self.tesseract.process(image)
|
||||
if self.engine_preference == "tesseract":
|
||||
return self.tesseract.process(preprocessed)
|
||||
|
||||
try:
|
||||
return self.paddle.process(original)
|
||||
except Exception as e:
|
||||
logger.warning("paddleocr_failed_fallback", error=str(e))
|
||||
return self.tesseract.process(preprocessed)
|
||||
|
||||
def process_file(self, image_path: Path) -> OCRResult:
|
||||
"""Load and process an image file."""
|
||||
|
||||
@@ -35,11 +35,13 @@ def _process_single_image(args: tuple[Path, dict]) -> ProcessedImage:
|
||||
try:
|
||||
image_hash = compute_image_hash(path)
|
||||
prep_start = time.time()
|
||||
image = preprocessor.preprocess(path)
|
||||
original = preprocessor.load_image(path)
|
||||
ocr_base = preprocessor.resize_if_needed(original)
|
||||
preprocessed = preprocessor.preprocess_image(original)
|
||||
preprocessing_time = time.time() - prep_start
|
||||
|
||||
ocr_start = time.time()
|
||||
result = ocr_engine.process(image)
|
||||
result = ocr_engine.process(ocr_base, preprocessed)
|
||||
ocr_time = time.time() - ocr_start
|
||||
|
||||
return ProcessedImage(
|
||||
@@ -87,7 +89,13 @@ class ParallelProcessor:
|
||||
unique_images = []
|
||||
seen = set()
|
||||
for img in all_images:
|
||||
if not img.is_file() or img.suffix.lower() not in (".png", ".jpg", ".jpeg", ".tiff", ".bmp"):
|
||||
if not img.is_file() or img.suffix.lower() not in (
|
||||
".png",
|
||||
".jpg",
|
||||
".jpeg",
|
||||
".tiff",
|
||||
".bmp",
|
||||
):
|
||||
continue
|
||||
if any(img.match(pattern) for pattern in settings.input.exclude_patterns):
|
||||
continue
|
||||
|
||||
@@ -27,7 +27,9 @@ class ImagePreprocessor:
|
||||
if max(height, width) <= self.max_dim:
|
||||
return image
|
||||
scale = self.max_dim / max(height, width)
|
||||
return cv2.resize(image, (int(width * scale), int(height * scale)), interpolation=cv2.INTER_AREA)
|
||||
return cv2.resize(
|
||||
image, (int(width * scale), int(height * scale)), interpolation=cv2.INTER_AREA
|
||||
)
|
||||
|
||||
def deskew(self, image: np.ndarray) -> np.ndarray:
|
||||
if not self.config.deskew:
|
||||
@@ -43,10 +45,16 @@ class ImagePreprocessor:
|
||||
return image
|
||||
height, width = image.shape[:2]
|
||||
matrix = cv2.getRotationMatrix2D((width // 2, height // 2), angle, 1.0)
|
||||
return cv2.warpAffine(image, matrix, (width, height), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)
|
||||
return cv2.warpAffine(
|
||||
image, matrix, (width, height), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE
|
||||
)
|
||||
|
||||
def denoise(self, image: np.ndarray) -> np.ndarray:
|
||||
return cv2.fastNlMeansDenoisingColored(image, None, 10, 10, 7, 21) if self.config.denoise else image
|
||||
return (
|
||||
cv2.fastNlMeansDenoisingColored(image, None, 10, 10, 7, 21)
|
||||
if self.config.denoise
|
||||
else image
|
||||
)
|
||||
|
||||
def apply_clahe(self, image: np.ndarray) -> np.ndarray:
|
||||
if not self.config.clahe:
|
||||
@@ -61,8 +69,12 @@ class ImagePreprocessor:
|
||||
return image
|
||||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||||
_, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
|
||||
horizontal = cv2.morphologyEx(binary, cv2.MORPH_OPEN, cv2.getStructuringElement(cv2.MORPH_RECT, (40, 1)))
|
||||
vertical = cv2.morphologyEx(binary, cv2.MORPH_OPEN, cv2.getStructuringElement(cv2.MORPH_RECT, (1, 40)))
|
||||
horizontal = cv2.morphologyEx(
|
||||
binary, cv2.MORPH_OPEN, cv2.getStructuringElement(cv2.MORPH_RECT, (40, 1))
|
||||
)
|
||||
vertical = cv2.morphologyEx(
|
||||
binary, cv2.MORPH_OPEN, cv2.getStructuringElement(cv2.MORPH_RECT, (1, 40))
|
||||
)
|
||||
mask = cv2.bitwise_or(horizontal, vertical)
|
||||
result = image.copy()
|
||||
result[mask > 0] = (255, 255, 255)
|
||||
@@ -72,7 +84,9 @@ class ImagePreprocessor:
|
||||
if not self.config.adaptive_threshold:
|
||||
return image
|
||||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||||
binary = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)
|
||||
binary = cv2.adaptiveThreshold(
|
||||
gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2
|
||||
)
|
||||
return cv2.cvtColor(binary, cv2.COLOR_GRAY2BGR)
|
||||
|
||||
def preprocess_image(self, image: np.ndarray) -> np.ndarray:
|
||||
|
||||
64
src/ocr_pipeline/onboarding.py
Normal file
64
src/ocr_pipeline/onboarding.py
Normal file
@@ -0,0 +1,64 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import platform
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
import yaml
|
||||
|
||||
from ocr_pipeline.config import DEFAULT_CONFIG_FILENAME, user_config_dir
|
||||
from ocr_pipeline.utils.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
CONFIG_HEADER = "# ocr-pipeline configuration (generated by `ocr-pipeline init`)\n"
|
||||
|
||||
|
||||
def default_screenshots_dir() -> Path:
|
||||
"""The directory where this OS puts screenshots by default."""
|
||||
system = platform.system().lower()
|
||||
if system == "darwin":
|
||||
return Path.home() / "Desktop"
|
||||
if system == "windows":
|
||||
return Path.home() / "Pictures" / "Screenshots"
|
||||
return Path.home() / "Pictures"
|
||||
|
||||
|
||||
def tesseract_path() -> str | None:
|
||||
"""Path to the tesseract executable, or None when it is not on PATH."""
|
||||
return shutil.which("tesseract")
|
||||
|
||||
|
||||
def tesseract_install_hint() -> str:
|
||||
"""One-line install instruction for the current platform."""
|
||||
system = platform.system().lower()
|
||||
if system == "darwin":
|
||||
return "brew install tesseract"
|
||||
if system == "windows":
|
||||
return "choco install tesseract # or https://github.com/UB-Mannheim/tesseract/wiki"
|
||||
return "sudo apt-get install tesseract-ocr # or your distro's package manager"
|
||||
|
||||
|
||||
def default_config(target_dir: Path, screenshots_dir: Path | None = None) -> dict:
|
||||
"""Minimal config content: safe Tesseract defaults, everything optional off."""
|
||||
shots = screenshots_dir or default_screenshots_dir()
|
||||
return {
|
||||
"input": {"paths": [str(shots)]},
|
||||
"ocr": {"engine": "tesseract"},
|
||||
"output": {"write_consolidated": True},
|
||||
}
|
||||
|
||||
|
||||
def write_default_config(target_dir: Path, screenshots_dir: Path | None = None) -> Path:
|
||||
"""Write a minimal config.yaml into target_dir and return its path."""
|
||||
target_dir.mkdir(parents=True, exist_ok=True)
|
||||
config_path = target_dir / DEFAULT_CONFIG_FILENAME
|
||||
body = yaml.safe_dump(default_config(target_dir, screenshots_dir), sort_keys=False)
|
||||
config_path.write_text(CONFIG_HEADER + body, encoding="utf-8")
|
||||
logger.info("config_written", path=str(config_path))
|
||||
return config_path
|
||||
|
||||
|
||||
def default_config_location(local: bool) -> Path:
|
||||
"""Where `init` writes config: cwd with --local, else the per-user config dir."""
|
||||
return Path.cwd() if local else user_config_dir()
|
||||
Binary file not shown.
Binary file not shown.
@@ -44,7 +44,11 @@ class MarkdownWriter:
|
||||
self._written_chunks: list[tuple[dict[str, Any], str]] = []
|
||||
|
||||
def start_run(self, timestamp: str, run_number: int) -> Path:
|
||||
"""Allocate one shared output directory for the entire pipeline run."""
|
||||
"""Allocate one shared output directory for the entire pipeline run.
|
||||
|
||||
The directory itself is created lazily on the first write, so a run
|
||||
where every image is skipped does not leave an empty run_NNN behind.
|
||||
"""
|
||||
self._written_chunks = []
|
||||
dt = datetime.datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
|
||||
if self.organize_by == "date_run":
|
||||
@@ -54,18 +58,20 @@ class MarkdownWriter:
|
||||
else:
|
||||
self._run_dir = None
|
||||
if self._run_dir is not None:
|
||||
self._run_dir.mkdir(parents=True, exist_ok=True)
|
||||
return self._run_dir
|
||||
return self.base_dir
|
||||
|
||||
def _get_output_dir(self, metadata: dict[str, Any]) -> Path:
|
||||
if self.organize_by == "source_dir":
|
||||
source_parent = Path(str(metadata.get("source_path", "unknown"))).parent
|
||||
safe_parent = "_".join(part for part in source_parent.parts if part not in (source_parent.anchor, "/"))
|
||||
safe_parent = "_".join(
|
||||
part for part in source_parent.parts if part not in (source_parent.anchor, "/")
|
||||
)
|
||||
output_dir = self.base_dir / "by_source" / (safe_parent[:100] or "unknown")
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
return output_dir
|
||||
if self._run_dir is not None:
|
||||
self._run_dir.mkdir(parents=True, exist_ok=True)
|
||||
return self._run_dir
|
||||
timestamp = metadata.get("timestamp")
|
||||
if not isinstance(timestamp, str) or not timestamp:
|
||||
@@ -84,16 +90,24 @@ class MarkdownWriter:
|
||||
) -> Path:
|
||||
output_dir = self._get_output_dir(metadata)
|
||||
source_path = Path(str(metadata.get("source_path", "unknown")))
|
||||
safe_name = "".join(c for c in source_path.stem if c.isalnum() or c in "-_")[:80] or "unknown"
|
||||
safe_name = (
|
||||
"".join(c for c in source_path.stem if c.isalnum() or c in "-_")[:80] or "unknown"
|
||||
)
|
||||
source_hash = str(metadata.get("source_hash", ""))[:12]
|
||||
chunk_idx = int(metadata.get("chunk_index", 0))
|
||||
total_chunks = int(metadata.get("total_chunks", 1))
|
||||
suffix = f"_{source_hash}" if source_hash else ""
|
||||
filename = f"{safe_name}{suffix}_chunk_{chunk_idx:03d}.md" if total_chunks > 1 else f"{safe_name}{suffix}.md"
|
||||
filename = (
|
||||
f"{safe_name}{suffix}_chunk_{chunk_idx:03d}.md"
|
||||
if total_chunks > 1
|
||||
else f"{safe_name}{suffix}.md"
|
||||
)
|
||||
output_path = output_dir / filename
|
||||
|
||||
frontmatter = self._build_frontmatter(metadata)
|
||||
body = self._build_body(content, source_path.name, figures or [], tables or [], entities, citations or [])
|
||||
body = self._build_body(
|
||||
content, source_path.name, figures or [], tables or [], entities, citations or []
|
||||
)
|
||||
self._written_chunks.append((dict(metadata), content))
|
||||
with output_path.open("w", encoding="utf-8") as handle:
|
||||
handle.write("---\n")
|
||||
@@ -112,7 +126,9 @@ class MarkdownWriter:
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
grouped: dict[str, list[tuple[dict[str, Any], str]]] = {}
|
||||
for metadata, content in self._written_chunks:
|
||||
grouped.setdefault(str(metadata.get("source_path", "unknown")), []).append((metadata, content))
|
||||
grouped.setdefault(str(metadata.get("source_path", "unknown")), []).append(
|
||||
(metadata, content)
|
||||
)
|
||||
|
||||
timestamp = datetime.datetime.now(datetime.UTC).isoformat()
|
||||
frontmatter = {
|
||||
@@ -125,7 +141,9 @@ class MarkdownWriter:
|
||||
parts = ["# Consolidated OCR Output", ""]
|
||||
for source_path, chunks in sorted(grouped.items()):
|
||||
parts.extend([f"## Source: {Path(source_path).name}", ""])
|
||||
for _metadata, content in sorted(chunks, key=lambda item: int(item[0].get("chunk_index", 0))):
|
||||
for _metadata, content in sorted(
|
||||
chunks, key=lambda item: int(item[0].get("chunk_index", 0))
|
||||
):
|
||||
parts.extend([content.strip(), ""])
|
||||
output_path = output_dir / settings.output.consolidated_filename
|
||||
with output_path.open("w", encoding="utf-8") as handle:
|
||||
@@ -133,7 +151,12 @@ class MarkdownWriter:
|
||||
yaml.safe_dump(frontmatter, handle, sort_keys=False, allow_unicode=True)
|
||||
handle.write("---\n\n")
|
||||
handle.write("\n".join(parts).rstrip() + "\n")
|
||||
logger.info("consolidated_markdown_written", path=str(output_path), sources=len(grouped), chunks=len(self._written_chunks))
|
||||
logger.info(
|
||||
"consolidated_markdown_written",
|
||||
path=str(output_path),
|
||||
sources=len(grouped),
|
||||
chunks=len(self._written_chunks),
|
||||
)
|
||||
return output_path
|
||||
|
||||
def _build_frontmatter(self, metadata: dict[str, Any]) -> dict[str, Any]:
|
||||
@@ -152,19 +175,38 @@ class MarkdownWriter:
|
||||
if figures:
|
||||
parts.extend(["## Figures Detected", ""])
|
||||
for index, figure in enumerate(figures, 1):
|
||||
parts.extend([f"### Figure {index}", f"- **Bounding Box**: {figure.bbox}", f"- **Confidence**: {figure.confidence:.2f}"])
|
||||
parts.extend(
|
||||
[
|
||||
f"### Figure {index}",
|
||||
f"- **Bounding Box**: {figure.bbox}",
|
||||
f"- **Confidence**: {figure.confidence:.2f}",
|
||||
]
|
||||
)
|
||||
if figure.caption:
|
||||
parts.append(f"- **Caption**: {figure.caption}")
|
||||
parts.append("")
|
||||
if tables:
|
||||
parts.extend(["## Tables Detected", ""])
|
||||
for index, table in enumerate(tables, 1):
|
||||
parts.extend([f"### Table {index}", f"- **Bounding Box**: {table.bbox}", f"- **Confidence**: {table.confidence:.2f}"])
|
||||
parts.extend(
|
||||
[
|
||||
f"### Table {index}",
|
||||
f"- **Bounding Box**: {table.bbox}",
|
||||
f"- **Confidence**: {table.confidence:.2f}",
|
||||
]
|
||||
)
|
||||
if table.markdown:
|
||||
parts.extend(["", "#### Table Content (Markdown)", "", table.markdown])
|
||||
parts.append("")
|
||||
if entities and entities.entities:
|
||||
parts.extend(["## Detected Entities", "", "| Entity | Type | Context |", "|--------|------|---------|"])
|
||||
parts.extend(
|
||||
[
|
||||
"## Detected Entities",
|
||||
"",
|
||||
"| Entity | Type | Context |",
|
||||
"|--------|------|---------|",
|
||||
]
|
||||
)
|
||||
for entity in entities.entities[:50]:
|
||||
context = entity.context[:80].replace("|", "\\|")
|
||||
if len(entity.context) > 80:
|
||||
@@ -172,7 +214,9 @@ class MarkdownWriter:
|
||||
parts.append(f"| {entity.text} | {entity.label} | {context} |")
|
||||
parts.append("")
|
||||
if citations:
|
||||
parts.extend(["## Citations Found", "", "| Type | Identifier |", "|------|------------|"])
|
||||
parts.extend(
|
||||
["## Citations Found", "", "| Type | Identifier |", "|------|------------|"]
|
||||
)
|
||||
for citation in citations:
|
||||
parts.append(f"| {citation.type.upper()} | `{citation.identifier}` |")
|
||||
parts.append("")
|
||||
|
||||
@@ -46,8 +46,10 @@ class OCRPipeline:
|
||||
return []
|
||||
|
||||
original_image = self.preprocessor.load_image(image_path)
|
||||
ocr_image = self.preprocessor.preprocess_image(original_image)
|
||||
ocr_result = self.engine.process(ocr_image)
|
||||
# PaddleOCR reads the (size-capped) original; Tesseract the binarized derivative.
|
||||
ocr_base = self.preprocessor.resize_if_needed(original_image)
|
||||
preprocessed = self.preprocessor.preprocess_image(original_image)
|
||||
ocr_result = self.engine.process(ocr_base, preprocessed)
|
||||
cleaned = clean_text(ocr_result.text)
|
||||
chunks = chunk_text(cleaned.text)
|
||||
entities = extract_entities(cleaned.text)
|
||||
@@ -60,16 +62,42 @@ class OCRPipeline:
|
||||
output_files: list[Path] = []
|
||||
for chunk in chunks:
|
||||
metadata = {
|
||||
"source_path": str(image_path), "source_hash": image_hash, "timestamp": timestamp,
|
||||
"ocr_engine": ocr_result.engine, "ocr_confidence_mean": ocr_result.confidence,
|
||||
"language": ocr_result.language, "entity_extraction_backend": entities.backend, "detected_entities": [entity.text for entity in entities.entities],
|
||||
"has_figures": bool(figures), "has_tables": bool(tables),
|
||||
"citations_found": [f"{citation.type}:{citation.identifier}" for citation in citations],
|
||||
"chunk_index": chunk.chunk_index, "total_chunks": len(chunks),
|
||||
"source_path": str(image_path),
|
||||
"source_hash": image_hash,
|
||||
"timestamp": timestamp,
|
||||
"ocr_engine": ocr_result.engine,
|
||||
"ocr_confidence_mean": ocr_result.confidence,
|
||||
"language": ocr_result.language,
|
||||
"entity_extraction_backend": entities.backend,
|
||||
"detected_entities": [entity.text for entity in entities.entities],
|
||||
"has_figures": bool(figures),
|
||||
"has_tables": bool(tables),
|
||||
"citations_found": [
|
||||
f"{citation.type}:{citation.identifier}" for citation in citations
|
||||
],
|
||||
"chunk_index": chunk.chunk_index,
|
||||
"total_chunks": len(chunks),
|
||||
}
|
||||
output_files.append(self.writer.write_chunk(content=chunk.content, metadata=metadata, figures=figures, tables=tables, entities=entities, citations=citations))
|
||||
output_files.append(
|
||||
self.writer.write_chunk(
|
||||
content=chunk.content,
|
||||
metadata=metadata,
|
||||
figures=figures,
|
||||
tables=tables,
|
||||
entities=entities,
|
||||
citations=citations,
|
||||
)
|
||||
)
|
||||
|
||||
self.db.mark_processed(file_hash=image_hash, path=str(image_path), timestamp=time.time(), engine=ocr_result.engine, confidence=ocr_result.confidence, chunks=len(chunks), run_id=self._run_id)
|
||||
self.db.mark_processed(
|
||||
file_hash=image_hash,
|
||||
path=str(image_path),
|
||||
timestamp=time.time(),
|
||||
engine=ocr_result.engine,
|
||||
confidence=ocr_result.confidence,
|
||||
chunks=len(chunks),
|
||||
run_id=self._run_id,
|
||||
)
|
||||
return output_files
|
||||
|
||||
def process_single(self, image_path: Path) -> list[Path]:
|
||||
@@ -108,7 +136,15 @@ class OCRPipeline:
|
||||
if consolidated is not None:
|
||||
result.output_files.append(consolidated)
|
||||
finally:
|
||||
self.db.complete_run(self._run_id, {"total": result.total_images, "successful": result.successful, "failed": result.failed, "chunks": result.chunks_created})
|
||||
self.db.complete_run(
|
||||
self._run_id,
|
||||
{
|
||||
"total": result.total_images,
|
||||
"successful": result.successful,
|
||||
"failed": result.failed,
|
||||
"chunks": result.chunks_created,
|
||||
},
|
||||
)
|
||||
self._run_id = None
|
||||
return result
|
||||
|
||||
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -74,4 +74,6 @@ def clean_text(text: str) -> CleanResult:
|
||||
text = re.sub(r"\n{3,}", "\n\n", text).strip()
|
||||
operations.append("normalize_whitespace")
|
||||
logger.debug("text_cleaned", original=original_length, cleaned=len(text), ops=operations)
|
||||
return CleanResult(text=text, original_length=original_length, cleaned_length=len(text), operations=operations)
|
||||
return CleanResult(
|
||||
text=text, original_length=original_length, cleaned_length=len(text), operations=operations
|
||||
)
|
||||
|
||||
@@ -3,8 +3,6 @@ from __future__ import annotations
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
|
||||
import spacy
|
||||
|
||||
from ocr_pipeline.config import settings
|
||||
from ocr_pipeline.utils.logging import get_logger
|
||||
|
||||
@@ -40,7 +38,17 @@ class ScientificEntityRecognizer:
|
||||
if self._initialized:
|
||||
return
|
||||
try:
|
||||
|
||||
import spacy
|
||||
except ImportError:
|
||||
spacy = None # type: ignore[assignment]
|
||||
if spacy is None:
|
||||
logger.info(
|
||||
"spacy_not_installed",
|
||||
hint="Install with: uv sync --extra scientific (regex entities still run)",
|
||||
)
|
||||
self._initialized = True
|
||||
return
|
||||
try:
|
||||
self._nlp = spacy.load(settings.entities.model)
|
||||
self._nlp.max_length = 2_000_000
|
||||
self._initialized = True
|
||||
@@ -48,10 +56,10 @@ class ScientificEntityRecognizer:
|
||||
logger.info("scispacy_model_loaded", model=settings.entities.model)
|
||||
except Exception as e:
|
||||
logger.warning("scispacy_load_failed", error=str(e), model=settings.entities.model)
|
||||
self._fallback_init()
|
||||
self._fallback_init(spacy)
|
||||
self._initialized = True
|
||||
|
||||
def _fallback_init(self):
|
||||
def _fallback_init(self, spacy):
|
||||
try:
|
||||
self._nlp = spacy.blank("en")
|
||||
self._nlp.add_pipe("sentencizer")
|
||||
@@ -90,7 +98,9 @@ class ScientificEntityRecognizer:
|
||||
)
|
||||
|
||||
entities = (
|
||||
self._merge_entities(entities, text) if settings.entities.merge_entities else entities
|
||||
self._merge_entities(entities, text)
|
||||
if settings.entities.merge_entities
|
||||
else entities
|
||||
)
|
||||
entities = self._deduplicate_entities(entities)
|
||||
|
||||
@@ -153,7 +163,6 @@ class RegexEntityRecognizer:
|
||||
],
|
||||
"CHEMICAL": [
|
||||
r"\b(?:DMSO|PBS|EDTA|Tris|HEPES|SDS|DTT|BME|NaCl|KCl|MgCl2|CaCl2|NaOH|HCl|H2SO4|HNO3|EtOH|MeOH|IPA|DMSO|DMF|DMA|THF|DCM|CHCl3|CH2Cl2|EtOAc|hexane|pentane|acetone|acetonitrile|water|H2O|buffer|media|serum|FBS|BSA|pen/strep|penicillin|streptomycin|trypsin|EDTA|collagenase|dispase|accutase)\b",
|
||||
|
||||
],
|
||||
"CONCENTRATION": [
|
||||
r"\b\d+(?:\.\d+)?\s*(?:[µumMpnc]?[Mm]|[µumMpnc]?[Mm]/[Ll]|[µumMpnc]?[Mm]\s*[Ll]?|[µumMpnc]?[gG]\s*/\s*[Ll]|[µumMpnc]?[gG]\s*/\s*[mM][lL]|[µumMpnc]?[gG]\s*/\s*[dL]|[µumMpnc]?[gG]\s*/\s*[mM][lL]|[µumMpnc]?[gG]\s*/\s*[dD][lL]|[µumMpnc]?[gG]\s*/\s*100\s*[mM][lL]|[µumMpnc]?[gG]\s*/\s*[kK][gG])\b",
|
||||
|
||||
131
src/ocr_pipeline/setup_steps.py
Normal file
131
src/ocr_pipeline/setup_steps.py
Normal file
@@ -0,0 +1,131 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import platform
|
||||
import shlex
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
|
||||
import typer
|
||||
from rich.console import Console
|
||||
|
||||
SCISPACY_MODEL_URL = (
|
||||
"https://s3-us-west-2.amazonaws.com/ai2-s2-scispacy/releases/v0.5.4/en_core_sci_lg-0.5.4.tar.gz"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SetupStep:
|
||||
"""One installation step. argv=None means the user must do it manually."""
|
||||
|
||||
label: str
|
||||
argv: tuple[str, ...] | None = None
|
||||
note: str | None = None
|
||||
|
||||
|
||||
def _tesseract_steps() -> list[SetupStep]:
|
||||
if shutil.which("tesseract"):
|
||||
return []
|
||||
system = platform.system().lower()
|
||||
if system == "darwin":
|
||||
if shutil.which("brew"):
|
||||
return [SetupStep("Install Tesseract via Homebrew", ("brew", "install", "tesseract"))]
|
||||
return [
|
||||
SetupStep(
|
||||
"Install Tesseract",
|
||||
None,
|
||||
"Install Homebrew first, then run: brew install tesseract",
|
||||
)
|
||||
]
|
||||
if system == "windows":
|
||||
if shutil.which("choco"):
|
||||
return [
|
||||
SetupStep("Install Tesseract via Chocolatey", ("choco", "install", "tesseract"))
|
||||
]
|
||||
return [
|
||||
SetupStep(
|
||||
"Install Tesseract",
|
||||
None,
|
||||
"Download the installer from https://github.com/UB-Mannheim/tesseract/wiki",
|
||||
)
|
||||
]
|
||||
if shutil.which("apt-get"):
|
||||
return [
|
||||
SetupStep(
|
||||
"Install Tesseract via apt", ("sudo", "apt-get", "install", "-y", "tesseract-ocr")
|
||||
)
|
||||
]
|
||||
return [
|
||||
SetupStep(
|
||||
"Install Tesseract", None, "Install the tesseract-ocr package for your distribution."
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def steps_for(target: str) -> list[SetupStep]:
|
||||
if target == "basic":
|
||||
steps = _tesseract_steps()
|
||||
steps.append(
|
||||
SetupStep("Verify the install", None, "Run: ocr-pipeline init && ocr-pipeline demo")
|
||||
)
|
||||
return steps
|
||||
if target == "full":
|
||||
steps: list[SetupStep] = []
|
||||
if shutil.which("uv"):
|
||||
steps.append(
|
||||
SetupStep("Install optional ML dependencies", ("uv", "sync", "--extra", "full"))
|
||||
)
|
||||
steps.append(
|
||||
SetupStep(
|
||||
"Download the scispaCy model (en_core_sci_lg)",
|
||||
("uv", "pip", "install", SCISPACY_MODEL_URL),
|
||||
"Requires a scispaCy-compatible spaCy; see the README if this step fails.",
|
||||
)
|
||||
)
|
||||
else:
|
||||
steps.append(
|
||||
SetupStep(
|
||||
"Install optional ML dependencies",
|
||||
(sys.executable, "-m", "pip", "install", "-e", ".[full]"),
|
||||
)
|
||||
)
|
||||
steps.append(
|
||||
SetupStep(
|
||||
"Download the scispaCy model (en_core_sci_lg)",
|
||||
(sys.executable, "-m", "pip", "install", SCISPACY_MODEL_URL),
|
||||
"Requires a scispaCy-compatible spaCy; see the README if this step fails.",
|
||||
)
|
||||
)
|
||||
steps.append(
|
||||
SetupStep(
|
||||
"Detectron2 for figure detection",
|
||||
None,
|
||||
"No universal wheel exists; install a build matching your torch/Python, see README.",
|
||||
)
|
||||
)
|
||||
return steps
|
||||
raise ValueError(f"Unknown setup target: {target}. Choose basic or full.")
|
||||
|
||||
|
||||
def run_steps(steps: list[SetupStep], *, assume_yes: bool, dry_run: bool, console: Console) -> int:
|
||||
"""Print and optionally execute each step. Returns the first failing exit code."""
|
||||
for index, step in enumerate(steps, 1):
|
||||
console.print(f"[cyan][{index}/{len(steps)}] {step.label}[/cyan]")
|
||||
if step.argv is None:
|
||||
if step.note:
|
||||
console.print(f" {step.note}")
|
||||
continue
|
||||
console.print(f" $ {shlex.join(step.argv)}")
|
||||
if dry_run:
|
||||
continue
|
||||
if step.note:
|
||||
console.print(f" {step.note}")
|
||||
if not assume_yes and not typer.confirm("Run this step?", default=True):
|
||||
console.print(" [dim]skipped[/dim]")
|
||||
continue
|
||||
returncode = subprocess.run(list(step.argv)).returncode
|
||||
if returncode != 0:
|
||||
console.print(f"[red]Step failed with exit code {returncode}.[/red]")
|
||||
return returncode
|
||||
return 0
|
||||
@@ -3,13 +3,11 @@ from __future__ import annotations
|
||||
from .db import Database, ProcessedFile, get_db
|
||||
from .logging import ProgressTracker, create_progress, get_logger, setup_logging
|
||||
|
||||
|
||||
__all__ = [
|
||||
"setup_logging",
|
||||
"get_logger",
|
||||
"create_progress",
|
||||
"ProgressTracker",
|
||||
|
||||
"get_db",
|
||||
"Database",
|
||||
"ProcessedFile",
|
||||
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -16,6 +16,7 @@ logger = get_logger(__name__)
|
||||
@dataclass(frozen=True)
|
||||
class ProcessedFile:
|
||||
"""Legacy value object retained for import compatibility."""
|
||||
|
||||
path: str
|
||||
hash: str
|
||||
timestamp: float
|
||||
@@ -75,38 +76,78 @@ class Database:
|
||||
def is_processed(self, file_hash: str, path: Path | str | None = None) -> bool:
|
||||
with self._conn() as conn:
|
||||
if path is None:
|
||||
row = conn.execute("SELECT 1 FROM processed_files WHERE hash = ?", (file_hash,)).fetchone()
|
||||
row = conn.execute(
|
||||
"SELECT 1 FROM processed_files WHERE hash = ?", (file_hash,)
|
||||
).fetchone()
|
||||
else:
|
||||
row = conn.execute("SELECT 1 FROM processed_files WHERE hash = ? AND path = ?", (file_hash, str(path))).fetchone()
|
||||
row = conn.execute(
|
||||
"SELECT 1 FROM processed_files WHERE hash = ? AND path = ?",
|
||||
(file_hash, str(path)),
|
||||
).fetchone()
|
||||
return row is not None
|
||||
|
||||
def get_processed(self, file_hash: str) -> dict[str, Any] | None:
|
||||
with self._conn() as conn:
|
||||
row = conn.execute("SELECT * FROM processed_files WHERE hash = ?", (file_hash,)).fetchone()
|
||||
row = conn.execute(
|
||||
"SELECT * FROM processed_files WHERE hash = ?", (file_hash,)
|
||||
).fetchone()
|
||||
return dict(row) if row else None
|
||||
|
||||
def mark_processed(self, *, file_hash: str, path: str, timestamp: float | str, engine: str = "", confidence: float = 0.0, chunks: int = 0, run_id: int | None = None) -> None:
|
||||
def mark_processed(
|
||||
self,
|
||||
*,
|
||||
file_hash: str,
|
||||
path: str,
|
||||
timestamp: float | str,
|
||||
engine: str = "",
|
||||
confidence: float = 0.0,
|
||||
chunks: int = 0,
|
||||
run_id: int | None = None,
|
||||
) -> None:
|
||||
with self._conn() as conn:
|
||||
conn.execute("""INSERT OR REPLACE INTO processed_files
|
||||
conn.execute(
|
||||
"""INSERT OR REPLACE INTO processed_files
|
||||
(hash, path, timestamp, ocr_engine, confidence, chunks_created, run_id)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)""", (file_hash, path, str(timestamp), engine, confidence, chunks, run_id))
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)""",
|
||||
(file_hash, path, str(timestamp), engine, confidence, chunks, run_id),
|
||||
)
|
||||
|
||||
def start_run(self, date: str | None = None) -> tuple[int, int]:
|
||||
date = date or time.strftime("%Y-%m-%d")
|
||||
with self._conn() as conn:
|
||||
run_number = conn.execute("SELECT COALESCE(MAX(run_number), 0) + 1 FROM runs WHERE date = ?", (date,)).fetchone()[0]
|
||||
cursor = conn.execute("INSERT INTO runs (date, run_number) VALUES (?, ?)", (date, run_number))
|
||||
run_number = conn.execute(
|
||||
"SELECT COALESCE(MAX(run_number), 0) + 1 FROM runs WHERE date = ?", (date,)
|
||||
).fetchone()[0]
|
||||
cursor = conn.execute(
|
||||
"INSERT INTO runs (date, run_number) VALUES (?, ?)", (date, run_number)
|
||||
)
|
||||
return int(cursor.lastrowid), int(run_number)
|
||||
|
||||
def complete_run(self, run_id: int, stats: dict[str, int]) -> None:
|
||||
with self._conn() as conn:
|
||||
conn.execute("""UPDATE runs SET total_images = ?, successful = ?, failed = ?, chunks_created = ?, completed_at = ? WHERE id = ?""", (stats.get("total", 0), stats.get("successful", 0), stats.get("failed", 0), stats.get("chunks", 0), time.time(), run_id))
|
||||
conn.execute(
|
||||
"""UPDATE runs SET total_images = ?, successful = ?, failed = ?, chunks_created = ?, completed_at = ? WHERE id = ?""",
|
||||
(
|
||||
stats.get("total", 0),
|
||||
stats.get("successful", 0),
|
||||
stats.get("failed", 0),
|
||||
stats.get("chunks", 0),
|
||||
time.time(),
|
||||
run_id,
|
||||
),
|
||||
)
|
||||
|
||||
def get_stats(self, days: int = 30) -> dict[str, Any]:
|
||||
cutoff = time.time() - days * 86400
|
||||
with self._conn() as conn:
|
||||
files = conn.execute("SELECT COUNT(*) AS total_processed, COALESCE(SUM(chunks_created), 0) AS total_chunks, AVG(confidence) AS avg_confidence FROM processed_files WHERE created_at >= ?", (cutoff,)).fetchone()
|
||||
runs = conn.execute("SELECT COUNT(*) AS total_runs, MAX(completed_at) AS latest_run FROM runs WHERE started_at >= ?", (cutoff,)).fetchone()
|
||||
files = conn.execute(
|
||||
"SELECT COUNT(*) AS total_processed, COALESCE(SUM(chunks_created), 0) AS total_chunks, AVG(confidence) AS avg_confidence FROM processed_files WHERE created_at >= ?",
|
||||
(cutoff,),
|
||||
).fetchone()
|
||||
runs = conn.execute(
|
||||
"SELECT COUNT(*) AS total_runs, MAX(completed_at) AS latest_run FROM runs WHERE started_at >= ?",
|
||||
(cutoff,),
|
||||
).fetchone()
|
||||
return {**dict(files), **dict(runs)}
|
||||
|
||||
|
||||
@@ -117,5 +158,6 @@ def get_db() -> Database:
|
||||
global _db_instance
|
||||
if _db_instance is None:
|
||||
from ocr_pipeline.config import settings
|
||||
|
||||
_db_instance = Database(settings.watch_db_path)
|
||||
return _db_instance
|
||||
|
||||
Binary file not shown.
Binary file not shown.
@@ -113,7 +113,15 @@ class Watcher:
|
||||
failed += 1
|
||||
|
||||
if self.run_id is not None:
|
||||
self.db.complete_run(self.run_id, {"total": processed + failed, "successful": processed, "failed": failed, "chunks": 0})
|
||||
self.db.complete_run(
|
||||
self.run_id,
|
||||
{
|
||||
"total": processed + failed,
|
||||
"successful": processed,
|
||||
"failed": failed,
|
||||
"chunks": 0,
|
||||
},
|
||||
)
|
||||
logger.info("watcher_stopped", run_id=self.run_id, processed=processed, failed=failed)
|
||||
|
||||
def _process_file(self, path: Path) -> None:
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
|
||||
|
||||
@@ -1,9 +1,8 @@
|
||||
import pytest
|
||||
from pathlib import Path
|
||||
|
||||
from ocr_pipeline.postprocess.clean import clean_text, CleanResult
|
||||
from ocr_pipeline.postprocess.chunk import chunk_text, TextChunk
|
||||
from ocr_pipeline.config import settings
|
||||
from ocr_pipeline.postprocess.chunk import TextChunk, chunk_text
|
||||
from ocr_pipeline.postprocess.clean import clean_text
|
||||
|
||||
|
||||
class TestTextCleaning:
|
||||
@@ -70,4 +69,4 @@ class TestConfig:
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
pytest.main([__file__, "-v"])
|
||||
|
||||
116
tests/test_startup.py
Normal file
116
tests/test_startup.py
Normal file
@@ -0,0 +1,116 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from ocr_pipeline import config as config_module
|
||||
from ocr_pipeline.cli import app
|
||||
|
||||
runner = CliRunner()
|
||||
|
||||
SRC_DIR = Path(__file__).parent.parent / "src"
|
||||
|
||||
|
||||
def _isolate_config(tmp_path: Path, monkeypatch) -> None:
|
||||
"""Pretend no config exists anywhere."""
|
||||
monkeypatch.delenv(config_module.CONFIG_ENV_VAR, raising=False)
|
||||
monkeypatch.chdir(tmp_path)
|
||||
monkeypatch.setattr(config_module, "user_config_dir", lambda: tmp_path / "no-such-dir")
|
||||
|
||||
|
||||
def test_find_config_file_prefers_env_var(tmp_path: Path, monkeypatch) -> None:
|
||||
env_cfg = tmp_path / "env.yaml"
|
||||
env_cfg.write_text("ocr:\n engine: tesseract\n", encoding="utf-8")
|
||||
cwd_dir = tmp_path / "cwd"
|
||||
cwd_dir.mkdir()
|
||||
(cwd_dir / "config.yaml").write_text("{}", encoding="utf-8")
|
||||
monkeypatch.setenv(config_module.CONFIG_ENV_VAR, str(env_cfg))
|
||||
monkeypatch.chdir(cwd_dir)
|
||||
assert config_module.find_config_file() == env_cfg
|
||||
|
||||
|
||||
def test_find_config_file_uses_cwd(tmp_path: Path, monkeypatch) -> None:
|
||||
monkeypatch.delenv(config_module.CONFIG_ENV_VAR, raising=False)
|
||||
(tmp_path / "config.yaml").write_text("{}", encoding="utf-8")
|
||||
monkeypatch.chdir(tmp_path)
|
||||
assert config_module.find_config_file() == tmp_path / "config.yaml"
|
||||
|
||||
|
||||
def test_find_config_file_none_when_unconfigured(tmp_path: Path, monkeypatch) -> None:
|
||||
_isolate_config(tmp_path, monkeypatch)
|
||||
assert config_module.find_config_file() is None
|
||||
|
||||
|
||||
def test_import_does_not_read_cwd_config(tmp_path: Path) -> None:
|
||||
(tmp_path / "config.yaml").write_text("{invalid yaml: [", encoding="utf-8")
|
||||
result = subprocess.run(
|
||||
[sys.executable, "-c", "import ocr_pipeline.config; print('ok')"],
|
||||
cwd=tmp_path,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
env={"PYTHONPATH": str(SRC_DIR), "PATH": "/usr/bin:/bin"},
|
||||
)
|
||||
assert result.returncode == 0
|
||||
assert "ok" in result.stdout
|
||||
|
||||
|
||||
def test_settings_proxy_delegates() -> None:
|
||||
from ocr_pipeline.config import settings
|
||||
|
||||
assert settings.ocr.engine in {"tesseract", "paddleocr", "auto"}
|
||||
assert repr(settings)
|
||||
|
||||
|
||||
def test_init_writes_local_config(tmp_path: Path, monkeypatch) -> None:
|
||||
_isolate_config(tmp_path, monkeypatch)
|
||||
result = runner.invoke(app, ["init", "--local"])
|
||||
assert result.exit_code == 0
|
||||
content = (tmp_path / "config.yaml").read_text(encoding="utf-8")
|
||||
assert "paths:" in content
|
||||
assert "tesseract" in content
|
||||
|
||||
|
||||
def test_init_refuses_to_overwrite(tmp_path: Path, monkeypatch) -> None:
|
||||
_isolate_config(tmp_path, monkeypatch)
|
||||
(tmp_path / "config.yaml").write_text("existing: true\n", encoding="utf-8")
|
||||
result = runner.invoke(app, ["init", "--local"])
|
||||
assert result.exit_code == 1
|
||||
assert (tmp_path / "config.yaml").read_text(encoding="utf-8") == "existing: true\n"
|
||||
|
||||
|
||||
def test_run_guard_without_config_or_input_dir(tmp_path: Path, monkeypatch) -> None:
|
||||
_isolate_config(tmp_path, monkeypatch)
|
||||
result = runner.invoke(app, ["run"])
|
||||
assert result.exit_code == 2
|
||||
assert "ocr-pipeline init" in result.output
|
||||
assert "ocr-pipeline demo" in result.output
|
||||
|
||||
|
||||
def test_watch_guard_without_config(tmp_path: Path, monkeypatch) -> None:
|
||||
_isolate_config(tmp_path, monkeypatch)
|
||||
result = runner.invoke(app, ["watch"])
|
||||
assert result.exit_code == 2
|
||||
|
||||
|
||||
def test_setup_dry_run_prints_steps_without_executing() -> None:
|
||||
result = runner.invoke(app, ["setup", "full", "--dry-run"])
|
||||
assert result.exit_code == 0
|
||||
assert "uv sync" in result.output or "pip install" in result.output
|
||||
|
||||
|
||||
def test_setup_rejects_unknown_target() -> None:
|
||||
result = runner.invoke(app, ["setup", "everything", "--dry-run"])
|
||||
assert result.exit_code == 2
|
||||
|
||||
|
||||
def test_demo_image_bundled_and_valid() -> None:
|
||||
from PIL import Image
|
||||
|
||||
from ocr_pipeline.demo import demo_image_path
|
||||
|
||||
with Image.open(demo_image_path()) as img:
|
||||
assert img.size[0] >= 800
|
||||
assert img.format == "PNG"
|
||||
Reference in New Issue
Block a user