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:
2026-07-19 22:14:54 +02:00
parent deab02610b
commit a0211155fa
59 changed files with 2426 additions and 629 deletions

View File

@@ -2,25 +2,24 @@ FROM python:3.11-slim
RUN apt-get update && apt-get install -y --no-install-recommends \ RUN apt-get update && apt-get install -y --no-install-recommends \
tesseract-ocr \ tesseract-ocr \
libtesseract-dev \
poppler-utils \
libgl1 \ libgl1 \
libglib2.0-0 \ libglib2.0-0 \
libsm6 \
libxext6 \
libxrender-dev \
libgomp1 \
&& rm -rf /var/lib/apt/lists/* && rm -rf /var/lib/apt/lists/*
WORKDIR /app WORKDIR /app
COPY pyproject.toml uv.lock ./ 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 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 COPY docker/config.docker.yaml ./config.docker.yaml
ENV OMP_NUM_THREADS=4
ENV PATH="/app/.venv/bin:$PATH" \
OCR_PIPELINE_CONFIG=/app/config.docker.yaml \
OMP_NUM_THREADS=4
ENTRYPOINT ["ocr-pipeline"] ENTRYPOINT ["ocr-pipeline"]

View File

@@ -11,61 +11,76 @@ Turns scientific screenshots into RAG-ready Markdown, with metadata attached.
## Quick start ## Quick start
The default install is small: Tesseract OCR only, no multi-GB ML downloads.
```bash ```bash
uv sync uv sync # install the core pipeline
pre-commit install uv run ocr-pipeline setup # installs Tesseract if missing (asks first)
uv run ocr-pipeline init # write a config for this machine
# Smoke test a small directory with the lightweight OCR path uv run ocr-pipeline demo # verify everything works on a sample screenshot
uv run ocr-pipeline run \ uv run ocr-pipeline run # process your screenshots
--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
``` ```
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 Use `--force` to reprocess files already recorded in the processing database. The
`--exclude` option is repeatable and prevents recursive scans from entering bundles `--exclude` option is repeatable and prevents recursive scans from entering bundles
such as macOS `.photoslibrary` directories. such as macOS `.photoslibrary` directories.
### Scientific NER environment ### Optional capabilities
`en_core_sci_lg` is distributed separately from scispaCy and is not installed by Heavy ML features are opt-in extras and degrade gracefully when absent (a warning
`uv sync`. The compatible scispaCy 0.5.4 release requires spaCy 3.7.x, so install it is logged and that step is skipped):
in a dedicated Python 3.11 environment rather than mixing it with the main Python
3.13 environment: | 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 ```bash
uv venv --python 3.11 .venv-scispacy uv run ocr-pipeline setup full # runs: uv sync --extra full + downloads en_core_sci_lg
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
``` ```
If the two environments are kept separate, invoke the pipeline with the explicit `en_core_sci_lg` is distributed separately from scispaCy. The compatible scispaCy
`.venv-scispacy/bin/ocr-pipeline` path; `uv run` selects the project's default `.venv`. 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 ## Docker
The image ships the Tesseract-only core with a ready-made config that reads from
`/data/screenshots` and writes to `/data/output`:
```bash ```bash
docker build -t ocr-pipeline . 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 # ML extras baked in:
docker run -v ~/Pictures:/data/screenshots -v ./data:/app/data ocr-pipeline watch docker build --build-arg EXTRAS="--extra full" -t ocr-pipeline:full .
``` ```
## Config ## Config
`ocr-pipeline init` writes a minimal working config; the example below shows every
knob for reference.
<details> <details>
<summary><code>config.yaml</code></summary> <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 ## 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/`: First run downloads what it needs, cached in `~/.cache/ocr_pipeline/`:
- PaddleOCR models (~200MB) - PaddleOCR models (~200MB)

View File

@@ -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: input:
paths: paths:
- "~/Pictures" - "~/Pictures"

15
docker/config.docker.yaml Normal file
View 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"

View File

@@ -1,5 +1,17 @@
from .ocr import get_screenshots, extract_text """Experimental txtai/litellm RAG interface built on the OCR pipeline.
from .embeddings import create_and_index
from .rag import retrieve, ask_wllm 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"] __all__ = ["get_screenshots", "extract_text", "create_and_index", "retrieve", "ask_wllm"]

View File

@@ -16,45 +16,22 @@ dependencies = [
"pyyaml>=6.0.1", "pyyaml>=6.0.1",
"structlog>=24.1.0", "structlog>=24.1.0",
"rich>=13.7.0", "rich>=13.7.0",
"tqdm>=4.66.0",
# Image processing # Image processing
"opencv-python-headless>=4.9.0", "opencv-python-headless>=4.9.0,<5.0.0",
"pillow>=10.2.0", "pillow>=10.2.0",
"numpy>=1.26.0,<2.0; python_version < '3.12'", "numpy>=1.26.0,<2.0; python_version < '3.12'",
"numpy>=1.26.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", "pytesseract>=0.3.10",
# Table detection # Text processing and notebook migration
"torch>=2.2.0",
"torchvision>=0.17.0",
"transformers>=4.38.0",
"timm>=1.0.0",
"layoutparser>=0.3.0",
# NLP
"langchain-text-splitters>=0.0.2", "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", "nbformat>=5.9.0",
"legacy-cgi>=2.6.2",
# Database
"sqlite-utils>=3.37.0",
# File watching # File watching
"watchdog>=3.0.0", "watchdog>=3.0.0",
# Utilities # Utilities
"python-slugify>=8.0.0",
"xxhash>=3.4.0",
"python-magic>=0.4.27",
"platformdirs>=4.2.0", "platformdirs>=4.2.0",
] ]
@@ -67,16 +44,43 @@ dev = [
"mypy>=1.8.0", "mypy>=1.8.0",
"pre-commit>=3.6.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 = [] grobid = []
full = [ full = [
"ocr-pipeline[dev]", "ocr-pipeline[paddle,tables,figures,scientific]",
] ]
[project.entry-points.console_scripts] [project.entry-points.console_scripts]
ocr-pipeline = "ocr_pipeline.cli:app" ocr-pipeline = "ocr_pipeline.cli:app"
[tool.setuptools.packages.find] [tool.setuptools.packages.find]
where = ["src"] where = ["src", "."]
include = ["ocr_pipeline*", "kg_ocr*"]
[tool.setuptools.package-data]
ocr_pipeline = ["assets/*.png"]
[tool.uv] [tool.uv]
dev-dependencies = [ dev-dependencies = [

View File

@@ -4,7 +4,7 @@ Version: 0.1.0
Summary: OCR pipeline for life science screenshots - RAG ready Summary: OCR pipeline for life science screenshots - RAG ready
Requires-Python: >=3.11 Requires-Python: >=3.11
Description-Content-Type: text/markdown 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>=2.6.0
Requires-Dist: pydantic-settings>=2.2.0 Requires-Dist: pydantic-settings>=2.2.0
Requires-Dist: pyyaml>=6.0.1 Requires-Dist: pyyaml>=6.0.1
@@ -13,20 +13,11 @@ Requires-Dist: rich>=13.7.0
Requires-Dist: tqdm>=4.66.0 Requires-Dist: tqdm>=4.66.0
Requires-Dist: opencv-python-headless>=4.9.0 Requires-Dist: opencv-python-headless>=4.9.0
Requires-Dist: pillow>=10.2.0 Requires-Dist: pillow>=10.2.0
Requires-Dist: numpy>=1.26.0 Requires-Dist: numpy<2.0,>=1.26.0; python_version < "3.12"
Requires-Dist: scikit-image>=0.22.0 Requires-Dist: numpy>=1.26.0; python_version >= "3.12"
Requires-Dist: paddleocr<3.0.0,>=2.7.0
Requires-Dist: paddlepaddle>=2.6.0
Requires-Dist: pytesseract>=0.3.10 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: langchain-text-splitters>=0.0.2
Requires-Dist: spacy>=3.7.0
Requires-Dist: nbformat>=5.9.0 Requires-Dist: nbformat>=5.9.0
Requires-Dist: legacy-cgi>=2.6.2
Requires-Dist: sqlite-utils>=3.37.0 Requires-Dist: sqlite-utils>=3.37.0
Requires-Dist: watchdog>=3.0.0 Requires-Dist: watchdog>=3.0.0
Requires-Dist: python-slugify>=8.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: mypy>=1.8.0; extra == "dev"
Requires-Dist: pre-commit>=3.6.0; extra == "dev" Requires-Dist: pre-commit>=3.6.0; extra == "dev"
Provides-Extra: grobid 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 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 # OCR Pipeline for Life Science Screenshots
Turns scientific screenshots into RAG-ready Markdown, with metadata attached. 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 - Preprocessing: deskew, denoise, CLAHE, line removal
- Figure/table detection - Optional figure/table detection
- Entity extraction with scispaCy - Optional scientific entity extraction with scispaCy
- Citation matching - Citation matching and embedding-ready chunking
- Chunking for embedding
## Quick start ## Quick start
```bash ```bash
uv sync # or: pip install -e . # The default basic profile is lightweight: Tesseract, cleanup, citations,
pre-commit install # chunking, Markdown, and no model downloads.
uv sync
uv run ocr-pipeline doctor
ocr-pipeline run # process screenshots # Smoke-test a small directory.
ocr-pipeline watch # watch a folder uv run ocr-pipeline run \
ocr-pipeline reprocess ocr_sc.ipynb # pull screenshots out of a notebook --input-dir ~/Pictures/test_screenshots \
ocr-pipeline status --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 ## Docker
```bash ```bash
@@ -86,9 +137,10 @@ input:
paths: ["~/Pictures", "/mnt/storage3/aman/screenshots"] paths: ["~/Pictures", "/mnt/storage3/aman/screenshots"]
patterns: ["SCR-*.png", "*.jpg", "*.jpeg", "*.tiff"] patterns: ["SCR-*.png", "*.jpg", "*.jpeg", "*.tiff"]
recursive: true recursive: true
exclude_patterns: ["*.photoslibrary/*"]
ocr: ocr:
engine: "paddleocr" # paddleocr | tesseract | auto engine: "tesseract" # paddleocr | tesseract | auto
languages: ["en", "latin"] languages: ["en", "latin"]
use_gpu: false use_gpu: false
preprocess: preprocess:
@@ -105,14 +157,14 @@ processing:
detectors: detectors:
figures: figures:
enabled: true enabled: false
confidence_threshold: 0.7 confidence_threshold: 0.7
tables: tables:
enabled: true enabled: false
confidence_threshold: 0.7 confidence_threshold: 0.7
entities: entities:
enabled: true enabled: false
model: "en_core_sci_lg" model: "en_core_sci_lg"
citations: citations:
@@ -126,6 +178,8 @@ chunking:
output: output:
base_directory: "./data/ocr_output" base_directory: "./data/ocr_output"
organize_by: "date_run" # date_run | source_dir | flat organize_by: "date_run" # date_run | source_dir | flat
write_consolidated: true
consolidated_filename: "all_ocr.md"
frontmatter: frontmatter:
- source_path - source_path
- source_hash - source_hash
@@ -134,6 +188,7 @@ output:
- ocr_confidence_mean - ocr_confidence_mean
- language - language
- detected_entities - detected_entities
- entity_extraction_backend
- has_figures - has_figures
- has_tables - has_tables
- citations_found - citations_found
@@ -200,15 +255,16 @@ Lane 4: BRCA1 KO + pBRCA1-C61G mutant
Anti-BRCA1 (1:1000), Anti-β-actin (1:5000) 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 ## 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. 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) - PaddleOCR models (~200MB)
- scispaCy `en_core_sci_lg` (~800MB) - 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. - **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. - **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 ## TODO

View File

@@ -4,6 +4,7 @@ setup.py
src/ocr_pipeline/__init__.py src/ocr_pipeline/__init__.py
src/ocr_pipeline/cli.py src/ocr_pipeline/cli.py
src/ocr_pipeline/config.py src/ocr_pipeline/config.py
src/ocr_pipeline/doctor.py
src/ocr_pipeline/pipeline.py src/ocr_pipeline/pipeline.py
src/ocr_pipeline/watch.py src/ocr_pipeline/watch.py
src/ocr_pipeline.egg-info/PKG-INFO src/ocr_pipeline.egg-info/PKG-INFO
@@ -39,3 +40,4 @@ tests/test_indexer.py
tests/test_ocr.py tests/test_ocr.py
tests/test_output.py tests/test_output.py
tests/test_postprocess.py tests/test_postprocess.py
tests/test_profiles.py

View File

@@ -1,4 +1,4 @@
typer[all]>=0.9.0 typer>=0.9.0
pydantic>=2.6.0 pydantic>=2.6.0
pydantic-settings>=2.2.0 pydantic-settings>=2.2.0
pyyaml>=6.0.1 pyyaml>=6.0.1
@@ -7,20 +7,9 @@ rich>=13.7.0
tqdm>=4.66.0 tqdm>=4.66.0
opencv-python-headless>=4.9.0 opencv-python-headless>=4.9.0
pillow>=10.2.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 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 langchain-text-splitters>=0.0.2
spacy>=3.7.0
nbformat>=5.9.0 nbformat>=5.9.0
legacy-cgi>=2.6.2
sqlite-utils>=3.37.0 sqlite-utils>=3.37.0
watchdog>=3.0.0 watchdog>=3.0.0
python-slugify>=8.0.0 python-slugify>=8.0.0
@@ -28,6 +17,12 @@ xxhash>=3.4.0
python-magic>=0.4.27 python-magic>=0.4.27
platformdirs>=4.2.0 platformdirs>=4.2.0
[:python_version < "3.12"]
numpy<2.0,>=1.26.0
[:python_version >= "3.12"]
numpy>=1.26.0
[dev] [dev]
pytest>=8.0.0 pytest>=8.0.0
pytest-cov>=4.1.0 pytest-cov>=4.1.0
@@ -36,7 +31,29 @@ ruff>=0.2.0
mypy>=1.8.0 mypy>=1.8.0
pre-commit>=3.6.0 pre-commit>=3.6.0
[figures]
layoutparser>=0.3.0
[full] [full]
ocr-pipeline[dev] ocr-pipeline[dev,figures,paddle,scientific,tables]
[grobid] [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

View File

Binary file not shown.

After

Width:  |  Height:  |  Size: 60 KiB

View File

@@ -13,4 +13,10 @@ from .regex_extractor import (
format_citations_for_markdown, format_citations_for_markdown,
) )
__all__ = ["Citation", "GrobidClient", "RegexCitationExtractor", "extract_citations", "format_citations_for_markdown"] __all__ = [
"Citation",
"GrobidClient",
"RegexCitationExtractor",
"extract_citations",
"format_citations_for_markdown",
]

View File

@@ -2,13 +2,22 @@ from __future__ import annotations
import sys import sys
from pathlib import Path from pathlib import Path
from typing import Literal, cast
import typer import typer
from rich.console import Console from rich.console import Console
from rich.table import Table 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.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.logging import get_logger, setup_logging
from ocr_pipeline.utils.migrate import migrate_notebook from ocr_pipeline.utils.migrate import migrate_notebook
@@ -21,6 +30,18 @@ console = Console()
logger = get_logger(__name__) 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() @app.callback()
def callback( def callback(
config: Path | None = typer.Option(None, "--config", "-c", help="Config file path"), config: Path | None = typer.Option(None, "--config", "-c", help="Config file path"),
@@ -31,23 +52,92 @@ def callback(
setup_logging(log_level, log_file) setup_logging(log_level, log_file)
if config: 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) 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() @app.command()
def run( def run(
input_dir: Path | None = typer.Option(None, "--input-dir", "-i", help="Input directory"), 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"), output_dir: Path | None = typer.Option(None, "--output-dir", "-o", help="Output directory"),
workers: int | None = typer.Option( workers: int | None = typer.Option(None, "--workers", "-w", help="Number of worker processes"),
None, "--workers", "-w", help="Number of worker processes"
),
engine: str | None = typer.Option( engine: str | None = typer.Option(
None, "--engine", "-e", help="OCR engine (paddleocr/tesseract/auto)" None, "--engine", "-e", help="OCR engine (paddleocr/tesseract/auto)"
), ),
exclude: list[str] = typer.Option([], "--exclude", help="Glob pattern to exclude; repeatable"), 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""" """Run OCR pipeline on all screenshots"""
_ensure_configured(input_dir)
if input_dir: if input_dir:
settings.input.paths = [str(input_dir)] settings.input.paths = [str(input_dir)]
if output_dir: if output_dir:
@@ -55,7 +145,12 @@ def run(
if workers: if workers:
settings.processing.workers = workers settings.processing.workers = workers
if engine: 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: if exclude:
settings.input.exclude_patterns = exclude settings.input.exclude_patterns = exclude
if force: if force:
@@ -70,6 +165,7 @@ def run(
table.add_row("Total Images", str(result.total_images)) table.add_row("Total Images", str(result.total_images))
table.add_row("Successful", str(result.successful)) 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("Failed", str(result.failed))
table.add_row("Chunks Created", str(result.chunks_created)) table.add_row("Chunks Created", str(result.chunks_created))
table.add_row("Output Files", str(len(result.output_files))) table.add_row("Output Files", str(len(result.output_files)))
@@ -88,6 +184,7 @@ def run(
@app.command() @app.command()
def watch(): def watch():
"""Watch for new screenshots and process automatically""" """Watch for new screenshots and process automatically"""
_ensure_configured(None)
pipeline = OCRPipeline() pipeline = OCRPipeline()
console.print("[green]Starting watch mode...[/green]") console.print("[green]Starting watch mode...[/green]")
console.print("Press Ctrl+C to stop") 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) 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() @app.command()
def config(): def config():
"""Show current configuration""" """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("[cyan]Current Configuration:[/cyan]")
console.print(settings.model_dump_json(indent=2)) console.print(settings.model_dump_json(indent=2))

View File

@@ -3,7 +3,7 @@ from __future__ import annotations
import os import os
import platform import platform
from pathlib import Path from pathlib import Path
from typing import Any, ClassVar, Literal from typing import Any, ClassVar, Literal, cast
import platformdirs import platformdirs
import yaml import yaml
@@ -12,6 +12,8 @@ from pydantic_settings import BaseSettings, SettingsConfigDict
APP_NAME = "ocr-pipeline" APP_NAME = "ocr-pipeline"
APP_AUTHOR = "aman" APP_AUTHOR = "aman"
CONFIG_ENV_VAR = "OCR_PIPELINE_CONFIG"
DEFAULT_CONFIG_FILENAME = "config.yaml"
def _get_platform_ignore_patterns() -> list[str]: def _get_platform_ignore_patterns() -> list[str]:
@@ -60,7 +62,7 @@ class PreprocessConfig(BaseSettings):
class OCRConfig(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"]) languages: list[str] = Field(default_factory=lambda: ["en", "latin"])
use_gpu: bool = False use_gpu: bool = False
use_angle_cls: bool = True use_angle_cls: bool = True
@@ -81,7 +83,7 @@ class ProcessingConfig(BaseSettings):
class DetectorConfig(BaseSettings): class DetectorConfig(BaseSettings):
enabled: bool = True enabled: bool = False
model: str = "" model: str = ""
confidence_threshold: float = Field(default=0.7, ge=0.0, le=1.0) confidence_threshold: float = Field(default=0.7, ge=0.0, le=1.0)
@@ -108,7 +110,7 @@ class DetectorsConfig(BaseSettings):
class EntitiesConfig(BaseSettings): class EntitiesConfig(BaseSettings):
enabled: bool = True enabled: bool = False
model: str = "en_core_sci_lg" model: str = "en_core_sci_lg"
types: list[str] = Field( types: list[str] = Field(
default_factory=lambda: [ default_factory=lambda: [
@@ -176,7 +178,9 @@ class WatchConfig(BaseSettings):
debounce_seconds: int = Field(default=5, ge=0) debounce_seconds: int = Field(default=5, ge=0)
ignore_patterns: list[str] = Field(default_factory=_get_platform_ignore_patterns) ignore_patterns: list[str] = Field(default_factory=_get_platform_ignore_patterns)
db_path: str = Field( 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) poll_interval: float = Field(default=1.0, gt=0)
@@ -189,8 +193,33 @@ class _BaseSettings(BaseSettings):
) )
def _load_yaml_config(path: str = "config.yaml") -> dict[str, Any]: def user_config_dir() -> Path:
"""Load YAML config file if it exists.""" """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) config_path = Path(path)
if not config_path.is_absolute(): if not config_path.is_absolute():
config_path = Path.cwd() / config_path config_path = Path.cwd() / config_path
@@ -198,6 +227,11 @@ def _load_yaml_config(path: str = "config.yaml") -> dict[str, Any]:
with open(config_path) as f: with open(config_path) as f:
return yaml.safe_load(f) or {} return yaml.safe_load(f) or {}
return {} 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]: 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) yaml_models = _build_settings_from_yaml(Settings._yaml_config)
super().__init__(**yaml_models, **kwargs) 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.""" """Reload settings from a YAML file."""
Settings._yaml_config = _load_yaml_config(path) Settings._yaml_config = _load_yaml_config(path)
# Re-initialize with new config # Re-initialize with new config
@@ -271,4 +305,34 @@ class Settings(_BaseSettings):
return Path(self.watch.db_path).expanduser().resolve() 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
View 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])

View File

@@ -33,9 +33,13 @@ class FigureDetector:
try: try:
import layoutparser as lp 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: 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( self._predictor = model_class(
config_path=settings.detectors.figures.model, config_path=settings.detectors.figures.model,
label_map={0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"}, label_map={0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"},
@@ -104,9 +108,13 @@ class CaptionExtractor:
try: try:
import layoutparser as lp 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: 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( self._predictor = model_class(
config_path=settings.detectors.captions.model, config_path=settings.detectors.captions.model,
label_map={0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"}, label_map={0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"},

View File

@@ -80,7 +80,12 @@ class TableDetector:
): ):
resized_box = [float(x) for x in box] resized_box = [float(x) for x in box]
scale_x, scale_y = original_width / 800, original_height / 800 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, y1, x2, y2 = map(int, box)
x1, x2 = max(0, x1), min(original_width, x2) x1, x2 = max(0, x1), min(original_width, x2)
y1, y2 = max(0, y1), min(original_height, y2) y1, y2 = max(0, y1), min(original_height, y2)

View File

@@ -68,7 +68,7 @@ class OCRResult:
class PaddleOCREngine: class PaddleOCREngine:
def __init__(self): def __init__(self):
self._ocr = None self._ocr: Any = None
self._initialized = False self._initialized = False
def _init(self): def _init(self):
@@ -104,6 +104,8 @@ class PaddleOCREngine:
def process(self, image: np.ndarray) -> OCRResult: def process(self, image: np.ndarray) -> OCRResult:
self._init() self._init()
if self._ocr is None:
raise RuntimeError("PaddleOCR failed to initialize")
start = time.time() start = time.time()
result = self._ocr.ocr(image, cls=True) result = self._ocr.ocr(image, cls=True)
@@ -211,24 +213,25 @@ class OCREngine:
self.tesseract = TesseractEngine() self.tesseract = TesseractEngine()
self.engine_preference = settings.ocr.engine self.engine_preference = settings.ocr.engine
def process(self, image: np.ndarray) -> OCRResult: def process(self, original: np.ndarray, preprocessed: np.ndarray | None = None) -> OCRResult:
"""Process a pre-loaded image array.""" """Process an image, giving each engine the variant it handles best.
if self.engine_preference == "paddleocr":
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
if self.engine_preference == "tesseract":
return self.tesseract.process(preprocessed)
try: try:
return self.paddle.process(image) return self.paddle.process(original)
except Exception as e: except Exception as e:
logger.warning("paddleocr_failed_fallback", error=str(e)) logger.warning("paddleocr_failed_fallback", error=str(e))
return self.tesseract.process(image) return self.tesseract.process(preprocessed)
elif self.engine_preference == "tesseract":
return self.tesseract.process(image)
else:
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_file(self, image_path: Path) -> OCRResult: def process_file(self, image_path: Path) -> OCRResult:
"""Load and process an image file.""" """Load and process an image file."""

View File

@@ -35,11 +35,13 @@ def _process_single_image(args: tuple[Path, dict]) -> ProcessedImage:
try: try:
image_hash = compute_image_hash(path) image_hash = compute_image_hash(path)
prep_start = time.time() 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 preprocessing_time = time.time() - prep_start
ocr_start = time.time() ocr_start = time.time()
result = ocr_engine.process(image) result = ocr_engine.process(ocr_base, preprocessed)
ocr_time = time.time() - ocr_start ocr_time = time.time() - ocr_start
return ProcessedImage( return ProcessedImage(
@@ -87,7 +89,13 @@ class ParallelProcessor:
unique_images = [] unique_images = []
seen = set() seen = set()
for img in all_images: 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 continue
if any(img.match(pattern) for pattern in settings.input.exclude_patterns): if any(img.match(pattern) for pattern in settings.input.exclude_patterns):
continue continue

View File

@@ -27,7 +27,9 @@ class ImagePreprocessor:
if max(height, width) <= self.max_dim: if max(height, width) <= self.max_dim:
return image return image
scale = self.max_dim / max(height, width) 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: def deskew(self, image: np.ndarray) -> np.ndarray:
if not self.config.deskew: if not self.config.deskew:
@@ -43,10 +45,16 @@ class ImagePreprocessor:
return image return image
height, width = image.shape[:2] height, width = image.shape[:2]
matrix = cv2.getRotationMatrix2D((width // 2, height // 2), angle, 1.0) 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: 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: def apply_clahe(self, image: np.ndarray) -> np.ndarray:
if not self.config.clahe: if not self.config.clahe:
@@ -61,8 +69,12 @@ class ImagePreprocessor:
return image return image
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
_, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU) _, 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))) horizontal = cv2.morphologyEx(
vertical = cv2.morphologyEx(binary, cv2.MORPH_OPEN, cv2.getStructuringElement(cv2.MORPH_RECT, (1, 40))) 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) mask = cv2.bitwise_or(horizontal, vertical)
result = image.copy() result = image.copy()
result[mask > 0] = (255, 255, 255) result[mask > 0] = (255, 255, 255)
@@ -72,7 +84,9 @@ class ImagePreprocessor:
if not self.config.adaptive_threshold: if not self.config.adaptive_threshold:
return image return image
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) 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) return cv2.cvtColor(binary, cv2.COLOR_GRAY2BGR)
def preprocess_image(self, image: np.ndarray) -> np.ndarray: def preprocess_image(self, image: np.ndarray) -> np.ndarray:

View 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()

View File

@@ -44,7 +44,11 @@ class MarkdownWriter:
self._written_chunks: list[tuple[dict[str, Any], str]] = [] self._written_chunks: list[tuple[dict[str, Any], str]] = []
def start_run(self, timestamp: str, run_number: int) -> Path: 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 = [] self._written_chunks = []
dt = datetime.datetime.fromisoformat(timestamp.replace("Z", "+00:00")) dt = datetime.datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
if self.organize_by == "date_run": if self.organize_by == "date_run":
@@ -54,18 +58,20 @@ class MarkdownWriter:
else: else:
self._run_dir = None self._run_dir = None
if self._run_dir is not None: if self._run_dir is not None:
self._run_dir.mkdir(parents=True, exist_ok=True)
return self._run_dir return self._run_dir
return self.base_dir return self.base_dir
def _get_output_dir(self, metadata: dict[str, Any]) -> Path: def _get_output_dir(self, metadata: dict[str, Any]) -> Path:
if self.organize_by == "source_dir": if self.organize_by == "source_dir":
source_parent = Path(str(metadata.get("source_path", "unknown"))).parent 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 = self.base_dir / "by_source" / (safe_parent[:100] or "unknown")
output_dir.mkdir(parents=True, exist_ok=True) output_dir.mkdir(parents=True, exist_ok=True)
return output_dir return output_dir
if self._run_dir is not None: if self._run_dir is not None:
self._run_dir.mkdir(parents=True, exist_ok=True)
return self._run_dir return self._run_dir
timestamp = metadata.get("timestamp") timestamp = metadata.get("timestamp")
if not isinstance(timestamp, str) or not timestamp: if not isinstance(timestamp, str) or not timestamp:
@@ -84,16 +90,24 @@ class MarkdownWriter:
) -> Path: ) -> Path:
output_dir = self._get_output_dir(metadata) output_dir = self._get_output_dir(metadata)
source_path = Path(str(metadata.get("source_path", "unknown"))) 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] source_hash = str(metadata.get("source_hash", ""))[:12]
chunk_idx = int(metadata.get("chunk_index", 0)) chunk_idx = int(metadata.get("chunk_index", 0))
total_chunks = int(metadata.get("total_chunks", 1)) total_chunks = int(metadata.get("total_chunks", 1))
suffix = f"_{source_hash}" if source_hash else "" 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 output_path = output_dir / filename
frontmatter = self._build_frontmatter(metadata) 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)) self._written_chunks.append((dict(metadata), content))
with output_path.open("w", encoding="utf-8") as handle: with output_path.open("w", encoding="utf-8") as handle:
handle.write("---\n") handle.write("---\n")
@@ -112,7 +126,9 @@ class MarkdownWriter:
output_dir.mkdir(parents=True, exist_ok=True) output_dir.mkdir(parents=True, exist_ok=True)
grouped: dict[str, list[tuple[dict[str, Any], str]]] = {} grouped: dict[str, list[tuple[dict[str, Any], str]]] = {}
for metadata, content in self._written_chunks: 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() timestamp = datetime.datetime.now(datetime.UTC).isoformat()
frontmatter = { frontmatter = {
@@ -125,7 +141,9 @@ class MarkdownWriter:
parts = ["# Consolidated OCR Output", ""] parts = ["# Consolidated OCR Output", ""]
for source_path, chunks in sorted(grouped.items()): for source_path, chunks in sorted(grouped.items()):
parts.extend([f"## Source: {Path(source_path).name}", ""]) 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(), ""]) parts.extend([content.strip(), ""])
output_path = output_dir / settings.output.consolidated_filename output_path = output_dir / settings.output.consolidated_filename
with output_path.open("w", encoding="utf-8") as handle: 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) yaml.safe_dump(frontmatter, handle, sort_keys=False, allow_unicode=True)
handle.write("---\n\n") handle.write("---\n\n")
handle.write("\n".join(parts).rstrip() + "\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 return output_path
def _build_frontmatter(self, metadata: dict[str, Any]) -> dict[str, Any]: def _build_frontmatter(self, metadata: dict[str, Any]) -> dict[str, Any]:
@@ -152,19 +175,38 @@ class MarkdownWriter:
if figures: if figures:
parts.extend(["## Figures Detected", ""]) parts.extend(["## Figures Detected", ""])
for index, figure in enumerate(figures, 1): 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: if figure.caption:
parts.append(f"- **Caption**: {figure.caption}") parts.append(f"- **Caption**: {figure.caption}")
parts.append("") parts.append("")
if tables: if tables:
parts.extend(["## Tables Detected", ""]) parts.extend(["## Tables Detected", ""])
for index, table in enumerate(tables, 1): 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: if table.markdown:
parts.extend(["", "#### Table Content (Markdown)", "", table.markdown]) parts.extend(["", "#### Table Content (Markdown)", "", table.markdown])
parts.append("") parts.append("")
if entities and entities.entities: 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]: for entity in entities.entities[:50]:
context = entity.context[:80].replace("|", "\\|") context = entity.context[:80].replace("|", "\\|")
if len(entity.context) > 80: if len(entity.context) > 80:
@@ -172,7 +214,9 @@ class MarkdownWriter:
parts.append(f"| {entity.text} | {entity.label} | {context} |") parts.append(f"| {entity.text} | {entity.label} | {context} |")
parts.append("") parts.append("")
if citations: if citations:
parts.extend(["## Citations Found", "", "| Type | Identifier |", "|------|------------|"]) parts.extend(
["## Citations Found", "", "| Type | Identifier |", "|------|------------|"]
)
for citation in citations: for citation in citations:
parts.append(f"| {citation.type.upper()} | `{citation.identifier}` |") parts.append(f"| {citation.type.upper()} | `{citation.identifier}` |")
parts.append("") parts.append("")

View File

@@ -46,8 +46,10 @@ class OCRPipeline:
return [] return []
original_image = self.preprocessor.load_image(image_path) original_image = self.preprocessor.load_image(image_path)
ocr_image = self.preprocessor.preprocess_image(original_image) # PaddleOCR reads the (size-capped) original; Tesseract the binarized derivative.
ocr_result = self.engine.process(ocr_image) 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) cleaned = clean_text(ocr_result.text)
chunks = chunk_text(cleaned.text) chunks = chunk_text(cleaned.text)
entities = extract_entities(cleaned.text) entities = extract_entities(cleaned.text)
@@ -60,16 +62,42 @@ class OCRPipeline:
output_files: list[Path] = [] output_files: list[Path] = []
for chunk in chunks: for chunk in chunks:
metadata = { metadata = {
"source_path": str(image_path), "source_hash": image_hash, "timestamp": timestamp, "source_path": str(image_path),
"ocr_engine": ocr_result.engine, "ocr_confidence_mean": ocr_result.confidence, "source_hash": image_hash,
"language": ocr_result.language, "entity_extraction_backend": entities.backend, "detected_entities": [entity.text for entity in entities.entities], "timestamp": timestamp,
"has_figures": bool(figures), "has_tables": bool(tables), "ocr_engine": ocr_result.engine,
"citations_found": [f"{citation.type}:{citation.identifier}" for citation in citations], "ocr_confidence_mean": ocr_result.confidence,
"chunk_index": chunk.chunk_index, "total_chunks": len(chunks), "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 return output_files
def process_single(self, image_path: Path) -> list[Path]: def process_single(self, image_path: Path) -> list[Path]:
@@ -108,7 +136,15 @@ class OCRPipeline:
if consolidated is not None: if consolidated is not None:
result.output_files.append(consolidated) result.output_files.append(consolidated)
finally: 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 self._run_id = None
return result return result

View File

@@ -74,4 +74,6 @@ def clean_text(text: str) -> CleanResult:
text = re.sub(r"\n{3,}", "\n\n", text).strip() text = re.sub(r"\n{3,}", "\n\n", text).strip()
operations.append("normalize_whitespace") operations.append("normalize_whitespace")
logger.debug("text_cleaned", original=original_length, cleaned=len(text), ops=operations) 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
)

View File

@@ -3,8 +3,6 @@ from __future__ import annotations
import re import re
from dataclasses import dataclass from dataclasses import dataclass
import spacy
from ocr_pipeline.config import settings from ocr_pipeline.config import settings
from ocr_pipeline.utils.logging import get_logger from ocr_pipeline.utils.logging import get_logger
@@ -40,7 +38,17 @@ class ScientificEntityRecognizer:
if self._initialized: if self._initialized:
return return
try: 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 = spacy.load(settings.entities.model)
self._nlp.max_length = 2_000_000 self._nlp.max_length = 2_000_000
self._initialized = True self._initialized = True
@@ -48,10 +56,10 @@ class ScientificEntityRecognizer:
logger.info("scispacy_model_loaded", model=settings.entities.model) logger.info("scispacy_model_loaded", model=settings.entities.model)
except Exception as e: except Exception as e:
logger.warning("scispacy_load_failed", error=str(e), model=settings.entities.model) logger.warning("scispacy_load_failed", error=str(e), model=settings.entities.model)
self._fallback_init() self._fallback_init(spacy)
self._initialized = True self._initialized = True
def _fallback_init(self): def _fallback_init(self, spacy):
try: try:
self._nlp = spacy.blank("en") self._nlp = spacy.blank("en")
self._nlp.add_pipe("sentencizer") self._nlp.add_pipe("sentencizer")
@@ -90,7 +98,9 @@ class ScientificEntityRecognizer:
) )
entities = ( 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) entities = self._deduplicate_entities(entities)
@@ -153,7 +163,6 @@ class RegexEntityRecognizer:
], ],
"CHEMICAL": [ "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", 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": [ "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", 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",

View 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

View File

@@ -3,13 +3,11 @@ from __future__ import annotations
from .db import Database, ProcessedFile, get_db from .db import Database, ProcessedFile, get_db
from .logging import ProgressTracker, create_progress, get_logger, setup_logging from .logging import ProgressTracker, create_progress, get_logger, setup_logging
__all__ = [ __all__ = [
"setup_logging", "setup_logging",
"get_logger", "get_logger",
"create_progress", "create_progress",
"ProgressTracker", "ProgressTracker",
"get_db", "get_db",
"Database", "Database",
"ProcessedFile", "ProcessedFile",

View File

@@ -16,6 +16,7 @@ logger = get_logger(__name__)
@dataclass(frozen=True) @dataclass(frozen=True)
class ProcessedFile: class ProcessedFile:
"""Legacy value object retained for import compatibility.""" """Legacy value object retained for import compatibility."""
path: str path: str
hash: str hash: str
timestamp: float timestamp: float
@@ -75,38 +76,78 @@ class Database:
def is_processed(self, file_hash: str, path: Path | str | None = None) -> bool: def is_processed(self, file_hash: str, path: Path | str | None = None) -> bool:
with self._conn() as conn: with self._conn() as conn:
if path is None: 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: 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 return row is not None
def get_processed(self, file_hash: str) -> dict[str, Any] | None: def get_processed(self, file_hash: str) -> dict[str, Any] | None:
with self._conn() as conn: 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 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: 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) (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]: def start_run(self, date: str | None = None) -> tuple[int, int]:
date = date or time.strftime("%Y-%m-%d") date = date or time.strftime("%Y-%m-%d")
with self._conn() as conn: with self._conn() as conn:
run_number = conn.execute("SELECT COALESCE(MAX(run_number), 0) + 1 FROM runs WHERE date = ?", (date,)).fetchone()[0] run_number = conn.execute(
cursor = conn.execute("INSERT INTO runs (date, run_number) VALUES (?, ?)", (date, run_number)) "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) return int(cursor.lastrowid), int(run_number)
def complete_run(self, run_id: int, stats: dict[str, int]) -> None: def complete_run(self, run_id: int, stats: dict[str, int]) -> None:
with self._conn() as conn: 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]: def get_stats(self, days: int = 30) -> dict[str, Any]:
cutoff = time.time() - days * 86400 cutoff = time.time() - days * 86400
with self._conn() as conn: 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() files = conn.execute(
runs = conn.execute("SELECT COUNT(*) AS total_runs, MAX(completed_at) AS latest_run FROM runs WHERE started_at >= ?", (cutoff,)).fetchone() "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)} return {**dict(files), **dict(runs)}
@@ -117,5 +158,6 @@ def get_db() -> Database:
global _db_instance global _db_instance
if _db_instance is None: if _db_instance is None:
from ocr_pipeline.config import settings from ocr_pipeline.config import settings
_db_instance = Database(settings.watch_db_path) _db_instance = Database(settings.watch_db_path)
return _db_instance return _db_instance

View File

@@ -113,7 +113,15 @@ class Watcher:
failed += 1 failed += 1
if self.run_id is not None: 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) logger.info("watcher_stopped", run_id=self.run_id, processed=processed, failed=failed)
def _process_file(self, path: Path) -> None: def _process_file(self, path: Path) -> None:

View File

@@ -1,9 +1,8 @@
import pytest 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.config import settings
from ocr_pipeline.postprocess.chunk import TextChunk, chunk_text
from ocr_pipeline.postprocess.clean import clean_text
class TestTextCleaning: class TestTextCleaning:

116
tests/test_startup.py Normal file
View 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"

1757
uv.lock generated

File diff suppressed because it is too large Load Diff