improve the OCR pipeline processing and outputs formatting
This commit is contained in:
@@ -14,8 +14,8 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
COPY pyproject.toml ./
|
||||
RUN pip install --no-cache-dir uv && uv sync --frozen
|
||||
COPY pyproject.toml uv.lock ./
|
||||
RUN pip install --no-cache-dir uv && uv sync --frozen --no-dev
|
||||
|
||||
COPY src ./src
|
||||
COPY config.yaml ./
|
||||
|
||||
13
README.md
13
README.md
@@ -154,7 +154,7 @@ 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
|
||||
|
||||
@@ -169,6 +169,17 @@ First run downloads what it needs, cached in `~/.cache/ocr_pipeline/`:
|
||||
- Table Transformer (~500MB)
|
||||
- LayoutParser PubLayNet (~300MB)
|
||||
|
||||
## Runtime requirements and health
|
||||
|
||||
The core pipeline can run with Tesseract alone. PaddleOCR, scientific NER, figure detection, and table detection are optional capabilities with heavyweight, platform-specific dependencies. The pipeline records the OCR engine actually used and the entity-extraction backend in generated frontmatter; inspect them after each run rather than assuming configured models loaded.
|
||||
|
||||
- **PaddleOCR:** the project pins the legacy 2.x API used by the pipeline. Install the locked environment with `uv sync`; a startup fallback to Tesseract is logged when Paddle cannot initialize.
|
||||
- **Scientific NER:** install a compatible `scispacy` distribution and the separately distributed `en_core_sci_lg` model before enabling production scientific NER. Without it, the pipeline uses its conservative regex fallback and marks the backend accordingly.
|
||||
- **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.
|
||||
|
||||
## TODO
|
||||
|
||||
- [ ] Build a knowledge graph from extracted entities/citations
|
||||
|
||||
@@ -8,6 +8,8 @@ input:
|
||||
- "*.tiff"
|
||||
- "*.bmp"
|
||||
recursive: true
|
||||
exclude_patterns:
|
||||
- "*.photoslibrary/*"
|
||||
|
||||
ocr:
|
||||
engine: "paddleocr"
|
||||
@@ -68,6 +70,8 @@ chunking:
|
||||
keep_separator: true
|
||||
|
||||
output:
|
||||
write_consolidated: true
|
||||
consolidated_filename: "all_ocr.md"
|
||||
# Uses platformdirs default: ~/.local/share/ocr-pipeline (Linux), ~/Library/Application Support/ocr-pipeline (macOS), %LOCALAPPDATA%\ocr-pipeline (Windows)
|
||||
format: "markdown"
|
||||
organize_by: "date_run"
|
||||
|
||||
@@ -1,3 +1,7 @@
|
||||
[build-system]
|
||||
requires = ["setuptools>=68"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "ocr-pipeline"
|
||||
version = "0.1.0"
|
||||
@@ -21,18 +25,22 @@ dependencies = [
|
||||
"scikit-image>=0.22.0",
|
||||
|
||||
# OCR engines
|
||||
"paddleocr>=2.7.0",
|
||||
"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
|
||||
"langchain-text-splitters>=0.0.2",
|
||||
"spacy>=3.7.0",
|
||||
"nbformat>=5.9.0",
|
||||
"legacy-cgi>=2.6.2",
|
||||
|
||||
# Database
|
||||
"sqlite-utils>=3.37.0",
|
||||
@@ -56,17 +64,17 @@ dev = [
|
||||
"mypy>=1.8.0",
|
||||
"pre-commit>=3.6.0",
|
||||
]
|
||||
grobid = [
|
||||
"grobid-client>=0.8.0",
|
||||
]
|
||||
grobid = []
|
||||
full = [
|
||||
"ocr-pipeline[dev]",
|
||||
"ocr-pipeline[grobid]",
|
||||
]
|
||||
|
||||
[project.entry-points.console_scripts]
|
||||
ocr-pipeline = "ocr_pipeline.cli:app"
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["src"]
|
||||
|
||||
[tool.uv]
|
||||
dev-dependencies = [
|
||||
"pytest>=8.0.0",
|
||||
|
||||
22
setup.py
22
setup.py
@@ -1,15 +1,9 @@
|
||||
from setuptools import setup, find_packages
|
||||
"""Legacy setuptools compatibility shim.
|
||||
|
||||
setup(
|
||||
name="kg_ocr",
|
||||
packages=find_packages(),
|
||||
python_requires=">=3.10",
|
||||
install_requires=[
|
||||
"pytesseract",
|
||||
"Pillow",
|
||||
"txtai",
|
||||
"sentence-transformers",
|
||||
"litellm",
|
||||
"python-dotenv"
|
||||
],
|
||||
)
|
||||
Project metadata, dependencies, and the src package layout are maintained solely in
|
||||
pyproject.toml. Install with `pip install .` or `uv sync`.
|
||||
"""
|
||||
|
||||
from setuptools import setup
|
||||
|
||||
setup()
|
||||
|
||||
@@ -15,19 +15,24 @@ 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>=2.7.0
|
||||
Requires-Dist: paddleocr<3.0.0,>=2.7.0
|
||||
Requires-Dist: paddlepaddle>=2.6.0
|
||||
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
|
||||
Requires-Dist: xxhash>=3.4.0
|
||||
Requires-Dist: python-magic>=0.4.27
|
||||
Requires-Dist: platformdirs>=4.2.0
|
||||
Provides-Extra: dev
|
||||
Requires-Dist: pytest>=8.0.0; extra == "dev"
|
||||
Requires-Dist: pytest-cov>=4.1.0; extra == "dev"
|
||||
@@ -36,57 +41,45 @@ 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
|
||||
Requires-Dist: grobid-client>=0.8.0; extra == "grobid"
|
||||
Provides-Extra: full
|
||||
Requires-Dist: ocr-pipeline[dev]; extra == "full"
|
||||
Requires-Dist: ocr-pipeline[grobid]; extra == "full"
|
||||
|
||||
# OCR Pipeline for Life Science Screenshots
|
||||
|
||||
A production-ready OCR pipeline that extracts text from scientific screenshots and converts them into RAG-ready Markdown files with rich metadata.
|
||||
Turns scientific screenshots into RAG-ready Markdown, with metadata attached.
|
||||
|
||||
## Features
|
||||
- OCR via PaddleOCR, falls back to Tesseract
|
||||
- Preprocessing: deskew, denoise, CLAHE, line removal
|
||||
- Figure/table detection
|
||||
- Entity extraction with scispaCy
|
||||
- Citation matching
|
||||
- Chunking for embedding
|
||||
|
||||
- **Multi-engine OCR**: PaddleOCR (primary) with Tesseract fallback
|
||||
- **Image preprocessing**: Deskewing, denoising, CLAHE contrast enhancement, line removal
|
||||
- **Parallel processing**: Multi-process worker pool for throughput
|
||||
- **Scientific entity recognition**: Genes, proteins, chemicals, species, diseases, cell lines via scispaCy
|
||||
- **Figure & table detection**: LayoutParser + Table Transformer
|
||||
- **Citation extraction**: DOI, PMID, arXiv, PMC, ISBN via regex + optional GROBID
|
||||
- **Semantic chunking**: LangChain recursive splitter optimized for scientific text
|
||||
- **RAG-ready output**: Markdown with YAML frontmatter (source, hash, timestamp, entities, confidence)
|
||||
- **Watch mode**: File system monitoring with SQLite persistence for incremental processing
|
||||
- **Notebook migration**: Reprocess screenshots from existing Jupyter notebooks
|
||||
- **Life-science focused**: Handles scientific notation, units, gene symbols, chemical formulas
|
||||
|
||||
## Quick Start
|
||||
## Quick start
|
||||
|
||||
```bash
|
||||
# Install with uv (recommended)
|
||||
uv sync
|
||||
|
||||
# Or with pip
|
||||
pip install -e .
|
||||
|
||||
# Install pre-commit hooks
|
||||
uv sync # or: pip install -e .
|
||||
pre-commit install
|
||||
|
||||
# Run on your screenshots
|
||||
ocr-pipeline run
|
||||
|
||||
# Watch for new screenshots
|
||||
ocr-pipeline watch
|
||||
|
||||
# Reprocess from notebook
|
||||
ocr-pipeline reprocess ocr_sc.ipynb
|
||||
|
||||
# Check status
|
||||
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
|
||||
```
|
||||
|
||||
## Configuration
|
||||
## Docker
|
||||
|
||||
Edit `config.yaml`:
|
||||
```bash
|
||||
docker build -t ocr-pipeline .
|
||||
|
||||
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
|
||||
```
|
||||
|
||||
## Config
|
||||
|
||||
<details>
|
||||
<summary><code>config.yaml</code></summary>
|
||||
|
||||
```yaml
|
||||
input:
|
||||
@@ -120,11 +113,11 @@ detectors:
|
||||
|
||||
entities:
|
||||
enabled: true
|
||||
model: "en_core_sci_lg" # scispaCy large model
|
||||
model: "en_core_sci_lg"
|
||||
|
||||
citations:
|
||||
regex_enabled: true
|
||||
grobid_enabled: false # Set true if running GROBID server
|
||||
grobid_enabled: false # set true if you're running a GROBID server
|
||||
|
||||
chunking:
|
||||
chunk_size: 1000
|
||||
@@ -153,9 +146,11 @@ watch:
|
||||
db_path: "./data/processed_files.db"
|
||||
```
|
||||
|
||||
## Output Format
|
||||
</details>
|
||||
|
||||
Each chunk produces a Markdown file with YAML frontmatter:
|
||||
## Output
|
||||
|
||||
Each chunk is a Markdown file with YAML frontmatter:
|
||||
|
||||
```markdown
|
||||
---
|
||||
@@ -178,9 +173,9 @@ total_chunks: 2
|
||||
## Figures
|
||||
|
||||
### Figure 1
|
||||
- **BBox**: [100, 200, 800, 600]
|
||||
- **Confidence**: 0.92
|
||||
- **Caption**: "Western blot showing BRCA1 expression..."
|
||||
- BBox: [100, 200, 800, 600]
|
||||
- Confidence: 0.92
|
||||
- Caption: "Western blot showing BRCA1 expression..."
|
||||
|
||||
## Detected Entities
|
||||
|
||||
@@ -205,72 +200,33 @@ Lane 4: BRCA1 KO + pBRCA1-C61G mutant
|
||||
Anti-BRCA1 (1:1000), Anti-β-actin (1:5000)
|
||||
```
|
||||
|
||||
## Directory Structure
|
||||
Files land in `data/ocr_output/<date>/run_NNN/`.
|
||||
|
||||
```
|
||||
data/ocr_output/
|
||||
├── 2025-01-15/
|
||||
│ ├── run_001/
|
||||
│ │ ├── SCR-20250115-gel_chunk_000.md
|
||||
│ │ └── SCR-20250115-gel_chunk_001.md
|
||||
│ └── run_002/
|
||||
│ └── ...
|
||||
└── 2025-01-16/
|
||||
└── run_001/
|
||||
└── ...
|
||||
```
|
||||
## Life science specifics
|
||||
|
||||
## Docker
|
||||
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.
|
||||
|
||||
```bash
|
||||
# Build
|
||||
docker build -t ocr-pipeline .
|
||||
## Models
|
||||
|
||||
# Run once
|
||||
docker run -v ~/Pictures:/data/screenshots -v ./data:/app/data ocr-pipeline run
|
||||
First run downloads what it needs, cached in `~/.cache/ocr_pipeline/`:
|
||||
|
||||
# Watch mode
|
||||
docker run -v ~/Pictures:/data/screenshots -v ./data:/app/data ocr-pipeline watch
|
||||
```
|
||||
|
||||
## Life Science Optimizations
|
||||
|
||||
| Feature | Implementation |
|
||||
|---------|----------------|
|
||||
| Gene/Protein names | scispaCy `en_core_sci_lg` NER |
|
||||
| Chemical formulas | Regex + unit normalization (µM, ng/mL, kb, etc.) |
|
||||
| Scientific notation | `1.5×10⁻³` → `1.5×10^-3` |
|
||||
| Gel/blot lanes | Figure detection + caption extraction |
|
||||
| Citations | DOI, PMID, arXiv, PMC, ISBN patterns |
|
||||
| Units | µM, ng/mL, kb/Mb/Gb, °C, ×g, etc. |
|
||||
|
||||
## Development
|
||||
|
||||
```bash
|
||||
# Install dev dependencies
|
||||
uv sync --dev
|
||||
|
||||
# Run tests
|
||||
pytest
|
||||
|
||||
# Lint
|
||||
ruff check .
|
||||
ruff format .
|
||||
|
||||
# Type check
|
||||
mypy src/
|
||||
```
|
||||
|
||||
## Model Downloads
|
||||
|
||||
First run downloads models automatically:
|
||||
- PaddleOCR detection/recognition models (~200MB)
|
||||
- PaddleOCR models (~200MB)
|
||||
- scispaCy `en_core_sci_lg` (~800MB)
|
||||
- Table Transformer (~500MB)
|
||||
- LayoutParser PubLayNet (~300MB)
|
||||
|
||||
Cache location: `~/.cache/ocr_pipeline/`
|
||||
## Runtime requirements and health
|
||||
|
||||
## License
|
||||
The core pipeline can run with Tesseract alone. PaddleOCR, scientific NER, figure detection, and table detection are optional capabilities with heavyweight, platform-specific dependencies. The pipeline records the OCR engine actually used and the entity-extraction backend in generated frontmatter; inspect them after each run rather than assuming configured models loaded.
|
||||
|
||||
- **PaddleOCR:** the project pins the legacy 2.x API used by the pipeline. Install the locked environment with `uv sync`; a startup fallback to Tesseract is logged when Paddle cannot initialize.
|
||||
- **Scientific NER:** install a compatible `scispacy` distribution and the separately distributed `en_core_sci_lg` model before enabling production scientific NER. Without it, the pipeline uses its conservative regex fallback and marks the backend accordingly.
|
||||
- **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.
|
||||
|
||||
## TODO
|
||||
|
||||
- [ ] Build a knowledge graph from extracted entities/citations
|
||||
|
||||
MIT
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
README.md
|
||||
pyproject.toml
|
||||
setup.py
|
||||
src/ocr_pipeline/__init__.py
|
||||
src/ocr_pipeline/cli.py
|
||||
src/ocr_pipeline/config.py
|
||||
@@ -20,6 +21,7 @@ src/ocr_pipeline/detectors/tables.py
|
||||
src/ocr_pipeline/ocr/__init__.py
|
||||
src/ocr_pipeline/ocr/engine.py
|
||||
src/ocr_pipeline/ocr/parallel.py
|
||||
src/ocr_pipeline/ocr/preprocess.py
|
||||
src/ocr_pipeline/output/__init__.py
|
||||
src/ocr_pipeline/output/markdown.py
|
||||
src/ocr_pipeline/postprocess/__init__.py
|
||||
@@ -32,4 +34,8 @@ src/ocr_pipeline/utils/logging.py
|
||||
src/ocr_pipeline/utils/migrate.py
|
||||
src/ocr_pipeline/watch/__init__.py
|
||||
src/ocr_pipeline/watch/watcher.py
|
||||
tests/test_graph.py
|
||||
tests/test_indexer.py
|
||||
tests/test_ocr.py
|
||||
tests/test_output.py
|
||||
tests/test_postprocess.py
|
||||
@@ -9,19 +9,24 @@ opencv-python-headless>=4.9.0
|
||||
pillow>=10.2.0
|
||||
numpy>=1.26.0
|
||||
scikit-image>=0.22.0
|
||||
paddleocr>=2.7.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
|
||||
xxhash>=3.4.0
|
||||
python-magic>=0.4.27
|
||||
platformdirs>=4.2.0
|
||||
|
||||
[dev]
|
||||
pytest>=8.0.0
|
||||
@@ -33,7 +38,5 @@ pre-commit>=3.6.0
|
||||
|
||||
[full]
|
||||
ocr-pipeline[dev]
|
||||
ocr-pipeline[grobid]
|
||||
|
||||
[grobid]
|
||||
grobid-client>=0.8.0
|
||||
|
||||
@@ -1,191 +1,16 @@
|
||||
from __future__ import annotations
|
||||
"""Backward-compatible citation exports.
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
The canonical implementation lives in ``regex_extractor``. Keeping this module as a
|
||||
thin re-export avoids two incompatible Citation models and removes an undeclared
|
||||
``requests`` dependency from the runtime path.
|
||||
"""
|
||||
|
||||
import requests
|
||||
|
||||
from ocr_pipeline.config import settings
|
||||
from ocr_pipeline.utils.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Citation:
|
||||
type: str
|
||||
value: str
|
||||
start: int
|
||||
end: int
|
||||
confidence: float
|
||||
metadata: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class RegexCitationExtractor:
|
||||
DOI_PATTERN = re.compile(r"\b(?:10\.\d{4,9}/[-._;()/:A-Z0-9]+)\b", re.IGNORECASE)
|
||||
PMID_PATTERN = re.compile(r"\b(?:PMID|PubMed|pubmed)[:/\s]?(\d{1,8})\b", re.IGNORECASE)
|
||||
ARXIV_PATTERN = re.compile(
|
||||
r"\b(?:arXiv|arxiv)[:/\s]?(\d{4}\.\d{4,5}(?:v\d+)?)\b", re.IGNORECASE
|
||||
)
|
||||
PMC_PATTERN = re.compile(r"\b(?:PMC|pmc)[:/\s]?(\d{1,8})\b", re.IGNORECASE)
|
||||
ISBN_PATTERN = re.compile(
|
||||
r"\b(?:ISBN[-:\s]?)?(?:97[89][-\s]?)?\d{1,5}[-\s]?\d{1,7}[-\s]?\d{1,7}[-\s]?[\dX]\b",
|
||||
re.IGNORECASE,
|
||||
from .regex_extractor import (
|
||||
Citation,
|
||||
GrobidClient,
|
||||
RegexCitationExtractor,
|
||||
extract_citations,
|
||||
format_citations_for_markdown,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def extract(cls, text: str) -> list[Citation]:
|
||||
citations = []
|
||||
|
||||
for match in cls.DOI_PATTERN.finditer(text):
|
||||
citations.append(
|
||||
Citation(
|
||||
type="DOI",
|
||||
value=match.group(),
|
||||
start=match.start(),
|
||||
end=match.end(),
|
||||
confidence=0.95,
|
||||
metadata={"url": f"https://doi.org/{match.group()}"},
|
||||
)
|
||||
)
|
||||
|
||||
for match in cls.PMID_PATTERN.finditer(text):
|
||||
citations.append(
|
||||
Citation(
|
||||
type="PMID",
|
||||
value=f"PMID:{match.group(1)}",
|
||||
start=match.start(),
|
||||
end=match.end(),
|
||||
confidence=0.9,
|
||||
metadata={"url": f"https://pubmed.ncbi.nlm.nih.gov/{match.group(1)}/"},
|
||||
)
|
||||
)
|
||||
|
||||
for match in cls.ARXIV_PATTERN.finditer(text):
|
||||
citations.append(
|
||||
Citation(
|
||||
type="arXiv",
|
||||
value=f"arXiv:{match.group(1)}",
|
||||
start=match.start(),
|
||||
end=match.end(),
|
||||
confidence=0.9,
|
||||
metadata={"url": f"https://arxiv.org/abs/{match.group(1)}"},
|
||||
)
|
||||
)
|
||||
|
||||
for match in cls.PMC_PATTERN.finditer(text):
|
||||
citations.append(
|
||||
Citation(
|
||||
type="PMC",
|
||||
value=f"PMC:{match.group(1)}",
|
||||
start=match.start(),
|
||||
end=match.end(),
|
||||
confidence=0.85,
|
||||
metadata={
|
||||
"url": f"https://www.ncbi.nlm.nih.gov/pmc/articles/PMC{match.group(1)}/"
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
for match in cls.ISBN_PATTERN.finditer(text):
|
||||
isbn = re.sub(r"[^\dX]", "", match.group(), flags=re.IGNORECASE)
|
||||
if len(isbn) in (10, 13):
|
||||
citations.append(
|
||||
Citation(
|
||||
type="ISBN",
|
||||
value=isbn,
|
||||
start=match.start(),
|
||||
end=match.end(),
|
||||
confidence=0.8,
|
||||
metadata={"url": f"https://isbnsearch.org/isbn/{isbn}"},
|
||||
)
|
||||
)
|
||||
|
||||
return citations
|
||||
|
||||
|
||||
class GrobidClient:
|
||||
def __init__(self):
|
||||
self.url = settings.citations.grobid_url
|
||||
self.timeout = settings.citations.grobid_timeout
|
||||
|
||||
def extract_references(self, text: str) -> list[Citation]:
|
||||
if not settings.citations.grobid_enabled:
|
||||
return []
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
f"{self.url}/api/processCitationList",
|
||||
data={"citations": text},
|
||||
headers={"Accept": "application/xml"},
|
||||
timeout=self.timeout,
|
||||
)
|
||||
if response.status_code != 200:
|
||||
logger.warning("grobid_request_failed", status=response.status_code)
|
||||
return []
|
||||
|
||||
return self._parse_grobid_xml(response.text)
|
||||
except Exception as e:
|
||||
logger.warning("grobid_extraction_failed", error=str(e))
|
||||
return []
|
||||
|
||||
def _parse_grobid_xml(self, xml_text: str) -> list[Citation]:
|
||||
citations = []
|
||||
try:
|
||||
import xml.etree.ElementTree as ET
|
||||
|
||||
root = ET.fromstring(xml_text)
|
||||
for cit in root.findall(".//{*}citation"):
|
||||
citation_text = ET.tostring(cit, encoding="unicode")
|
||||
doi_match = self.DOI_PATTERN.search(citation_text)
|
||||
if doi_match:
|
||||
citations.append(
|
||||
Citation(
|
||||
type="DOI",
|
||||
value=doi_match.group(),
|
||||
start=0,
|
||||
end=len(citation_text),
|
||||
confidence=0.99,
|
||||
metadata={"source": "grobid", "full_citation": citation_text},
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return citations
|
||||
|
||||
|
||||
def extract_citations(text: str) -> list[Citation]:
|
||||
citations = []
|
||||
|
||||
if settings.citations.regex_enabled:
|
||||
citations.extend(RegexCitationExtractor.extract(text))
|
||||
|
||||
if settings.citations.grobid_enabled:
|
||||
grobid = GrobidClient()
|
||||
citations.extend(grobid.extract_references(text))
|
||||
|
||||
seen = set()
|
||||
unique = []
|
||||
for cit in citations:
|
||||
key = (cit.type, cit.value)
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
unique.append(cit)
|
||||
|
||||
return unique
|
||||
|
||||
|
||||
def format_citations_for_markdown(citations: list[Citation]) -> str:
|
||||
if not citations:
|
||||
return "No citations found."
|
||||
|
||||
lines = ["## Citations", ""]
|
||||
for cit in citations:
|
||||
url = cit.metadata.get("url", "") if cit.metadata else ""
|
||||
if url:
|
||||
lines.append(f"- [{cit.type}: {cit.value}]({url})")
|
||||
else:
|
||||
lines.append(f"- {cit.type}: {cit.value}")
|
||||
|
||||
return "\n".join(lines)
|
||||
__all__ = ["Citation", "GrobidClient", "RegexCitationExtractor", "extract_citations", "format_citations_for_markdown"]
|
||||
|
||||
@@ -2,8 +2,8 @@ from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
|
||||
import httpx
|
||||
from urllib.parse import urlencode
|
||||
from urllib.request import Request, urlopen
|
||||
|
||||
from ocr_pipeline.config import settings
|
||||
from ocr_pipeline.utils.logging import get_logger
|
||||
@@ -26,7 +26,7 @@ class RegexCitationExtractor:
|
||||
re.IGNORECASE,
|
||||
)
|
||||
PMID_PATTERN = re.compile(
|
||||
r"\b(?:PMID|PubMed\s*ID|PubMed)\s*[:\-]?\s*(\d{8})\b",
|
||||
r"\b(?:PMID|PubMed\s*ID|PubMed)\s*[:\-]?\s*(\d{1,8})\b",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
ARXIV_PATTERN = re.compile(
|
||||
@@ -49,7 +49,7 @@ class RegexCitationExtractor:
|
||||
Citation(
|
||||
raw=match.group(),
|
||||
type="doi",
|
||||
identifier=match.group().lower(),
|
||||
identifier=match.group().rstrip(".,;:)]}").lower(),
|
||||
confidence=0.95,
|
||||
context=text[ctx_start:ctx_end].strip(),
|
||||
)
|
||||
@@ -108,13 +108,14 @@ class GrobidClient:
|
||||
return []
|
||||
|
||||
try:
|
||||
with httpx.Client(timeout=self.timeout) as client:
|
||||
response = client.post(
|
||||
payload = urlencode({"text": text}).encode("utf-8")
|
||||
request = Request(
|
||||
f"{self.url}/api/processCitationText",
|
||||
data={"text": text},
|
||||
data=payload,
|
||||
headers={"Content-Type": "application/x-www-form-urlencoded"},
|
||||
)
|
||||
response.raise_for_status()
|
||||
return self._parse_grobid_xml(response.text)
|
||||
with urlopen(request, timeout=self.timeout) as response:
|
||||
return self._parse_grobid_xml(response.read().decode("utf-8"))
|
||||
except Exception as e:
|
||||
logger.warning("grobid_request_failed", error=str(e))
|
||||
return []
|
||||
|
||||
@@ -44,6 +44,8 @@ def run(
|
||||
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"),
|
||||
):
|
||||
"""Run OCR pipeline on all screenshots"""
|
||||
if input_dir:
|
||||
@@ -54,6 +56,10 @@ def run(
|
||||
settings.processing.workers = workers
|
||||
if engine:
|
||||
settings.ocr.engine = engine
|
||||
if exclude:
|
||||
settings.input.exclude_patterns = exclude
|
||||
if force:
|
||||
settings.processing.skip_existing = False
|
||||
|
||||
pipeline = OCRPipeline()
|
||||
result = pipeline.process_all()
|
||||
@@ -136,7 +142,7 @@ def status():
|
||||
def stats(days: int = typer.Option(30, "--days", "-d", help="Number of days to analyze")):
|
||||
"""Show processing statistics"""
|
||||
pipeline = OCRPipeline()
|
||||
stats = pipeline.get_stats()
|
||||
stats = pipeline.get_stats(days)
|
||||
|
||||
table = Table(title=f"Statistics (last {days} days)")
|
||||
table.add_column("Metric", style="cyan")
|
||||
|
||||
@@ -3,14 +3,13 @@ from __future__ import annotations
|
||||
import os
|
||||
import platform
|
||||
from pathlib import Path
|
||||
from typing import Any, ClassVar
|
||||
from typing import Any, ClassVar, Literal
|
||||
|
||||
import platformdirs
|
||||
import yaml
|
||||
from pydantic import Field, field_validator
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
|
||||
APP_NAME = "ocr-pipeline"
|
||||
APP_AUTHOR = "aman"
|
||||
|
||||
@@ -43,6 +42,7 @@ class InputConfig(BaseSettings):
|
||||
default_factory=lambda: ["SCR-*.png", "*.jpg", "*.jpeg", "*.tiff", "*.bmp"]
|
||||
)
|
||||
recursive: bool = True
|
||||
exclude_patterns: list[str] = Field(default_factory=lambda: ["*.photoslibrary/*"])
|
||||
|
||||
@field_validator("paths", mode="before")
|
||||
@classmethod
|
||||
@@ -60,7 +60,7 @@ class PreprocessConfig(BaseSettings):
|
||||
|
||||
|
||||
class OCRConfig(BaseSettings):
|
||||
engine: str = "paddleocr" # paddleocr | tesseract | auto
|
||||
engine: Literal["paddleocr", "tesseract", "auto"] = "paddleocr"
|
||||
languages: list[str] = Field(default_factory=lambda: ["en", "latin"])
|
||||
use_gpu: bool = False
|
||||
use_angle_cls: bool = True
|
||||
@@ -73,32 +73,32 @@ class OCRConfig(BaseSettings):
|
||||
|
||||
|
||||
class ProcessingConfig(BaseSettings):
|
||||
workers: int = 4
|
||||
batch_size: int = 10
|
||||
retry_attempts: int = 2
|
||||
timeout_per_image: int = 120
|
||||
workers: int = Field(default=4, ge=1)
|
||||
batch_size: int = Field(default=10, ge=1)
|
||||
retry_attempts: int = Field(default=2, ge=0)
|
||||
timeout_per_image: int = Field(default=120, ge=1)
|
||||
skip_existing: bool = True
|
||||
|
||||
|
||||
class DetectorConfig(BaseSettings):
|
||||
enabled: bool = True
|
||||
model: str = ""
|
||||
confidence_threshold: float = 0.7
|
||||
confidence_threshold: float = Field(default=0.7, ge=0.0, le=1.0)
|
||||
|
||||
|
||||
class FigureDetectorConfig(DetectorConfig):
|
||||
model: str = "lp://PubLayNet/faster_rcnn_R_50_FPN_3x/config"
|
||||
confidence_threshold: float = 0.7
|
||||
confidence_threshold: float = Field(default=0.7, ge=0.0, le=1.0)
|
||||
|
||||
|
||||
class TableDetectorConfig(DetectorConfig):
|
||||
model: str = "microsoft/table-transformer-detection"
|
||||
confidence_threshold: float = 0.7
|
||||
confidence_threshold: float = Field(default=0.7, ge=0.0, le=1.0)
|
||||
|
||||
|
||||
class CaptionDetectorConfig(DetectorConfig):
|
||||
model: str = "lp://PubLayNet/faster_rcnn_R_50_FPN_3x/config"
|
||||
confidence_threshold: float = 0.5
|
||||
confidence_threshold: float = Field(default=0.5, ge=0.0, le=1.0)
|
||||
|
||||
|
||||
class DetectorsConfig(BaseSettings):
|
||||
@@ -149,6 +149,7 @@ class FrontmatterConfig(BaseSettings):
|
||||
"ocr_confidence_mean",
|
||||
"language",
|
||||
"detected_entities",
|
||||
"entity_extraction_backend",
|
||||
"has_figures",
|
||||
"has_tables",
|
||||
"citations_found",
|
||||
@@ -160,22 +161,24 @@ class FrontmatterConfig(BaseSettings):
|
||||
|
||||
|
||||
class OutputConfig(BaseSettings):
|
||||
write_consolidated: bool = False
|
||||
consolidated_filename: str = "all_ocr.md"
|
||||
base_directory: str = Field(
|
||||
default_factory=lambda: str(platformdirs.user_data_dir(APP_NAME, APP_AUTHOR))
|
||||
)
|
||||
format: str = "markdown"
|
||||
organize_by: str = "date_run" # date_run | source_dir | flat
|
||||
organize_by: Literal["date_run", "source_dir", "flat"] = "date_run"
|
||||
frontmatter: FrontmatterConfig = Field(default_factory=FrontmatterConfig)
|
||||
|
||||
|
||||
class WatchConfig(BaseSettings):
|
||||
enabled: bool = True
|
||||
debounce_seconds: int = 5
|
||||
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")
|
||||
)
|
||||
poll_interval: float = 1.0
|
||||
poll_interval: float = Field(default=1.0, gt=0)
|
||||
|
||||
|
||||
class _BaseSettings(BaseSettings):
|
||||
|
||||
@@ -33,7 +33,10 @@ class FigureDetector:
|
||||
try:
|
||||
import layoutparser as lp
|
||||
|
||||
self._predictor = lp.Detectron2LayoutModel(
|
||||
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")
|
||||
self._predictor = model_class(
|
||||
config_path=settings.detectors.figures.model,
|
||||
label_map={0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"},
|
||||
extra_config=[
|
||||
@@ -67,7 +70,10 @@ class FigureDetector:
|
||||
block.block.x_2,
|
||||
block.block.y_2,
|
||||
)
|
||||
crop = image[int(y1) : int(y2), int(x1) : int(x2)]
|
||||
height, width = image.shape[:2]
|
||||
x1_i, x2_i = max(0, int(x1)), min(width, int(x2))
|
||||
y1_i, y2_i = max(0, int(y1)), min(height, int(y2))
|
||||
crop = image[y1_i:y2_i, x1_i:x2_i]
|
||||
figures.append(
|
||||
DetectedFigure(
|
||||
bbox=[float(x1), float(y1), float(x2), float(y2)],
|
||||
@@ -98,7 +104,10 @@ class CaptionExtractor:
|
||||
try:
|
||||
import layoutparser as lp
|
||||
|
||||
self._predictor = lp.Detectron2LayoutModel(
|
||||
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")
|
||||
self._predictor = model_class(
|
||||
config_path=settings.detectors.captions.model,
|
||||
label_map={0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"},
|
||||
extra_config=[
|
||||
@@ -156,11 +165,14 @@ class CaptionExtractor:
|
||||
return figures
|
||||
|
||||
|
||||
def detect_figures(image: np.ndarray) -> list[DetectedFigure]:
|
||||
detector = FigureDetector()
|
||||
figures = detector.detect(image)
|
||||
_figure_detector = FigureDetector()
|
||||
_caption_extractor = CaptionExtractor()
|
||||
|
||||
captioner = CaptionExtractor()
|
||||
|
||||
def detect_figures(image: np.ndarray) -> list[DetectedFigure]:
|
||||
figures = _figure_detector.detect(image)
|
||||
|
||||
captioner = _caption_extractor
|
||||
figures = captioner.extract_captions(image, figures)
|
||||
|
||||
return figures
|
||||
|
||||
@@ -59,6 +59,7 @@ class TableDetector:
|
||||
try:
|
||||
import torch
|
||||
|
||||
original_height, original_width = image.shape[:2]
|
||||
pil_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||||
pil_image = cv2.resize(pil_image, (800, 800))
|
||||
|
||||
@@ -77,8 +78,12 @@ class TableDetector:
|
||||
for score, _label, box in zip(
|
||||
results["scores"], results["labels"], results["boxes"], strict=True
|
||||
):
|
||||
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
|
||||
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)
|
||||
crop = image[y1:y2, x1:x2]
|
||||
tables.append(
|
||||
DetectedTable(
|
||||
@@ -122,6 +127,7 @@ class TableStructureRecognizer:
|
||||
self._initialized = True
|
||||
|
||||
def recognize(self, table_crop: np.ndarray) -> tuple[list[dict[str, Any]] | None, str | None]:
|
||||
self._init()
|
||||
if self._model is None or self._processor is None:
|
||||
return None, None
|
||||
|
||||
@@ -150,8 +156,9 @@ class TableStructureRecognizer:
|
||||
}
|
||||
)
|
||||
|
||||
markdown = self._structure_to_markdown(structure, table_crop.shape)
|
||||
return structure, markdown
|
||||
# Structure models locate rows/cells but do not provide cell text. Do not
|
||||
# emit labels as fabricated table content; OCR must be associated separately.
|
||||
return structure, None
|
||||
except Exception as e:
|
||||
logger.warning("table_structure_recognition_failed", error=str(e))
|
||||
return None, None
|
||||
@@ -185,17 +192,19 @@ class TableStructureRecognizer:
|
||||
return None
|
||||
|
||||
|
||||
_table_detector = TableDetector()
|
||||
_table_structure_recognizer = TableStructureRecognizer()
|
||||
|
||||
|
||||
def detect_tables(image: np.ndarray) -> list[DetectedTable]:
|
||||
detector = TableDetector()
|
||||
tables = detector.detect(image)
|
||||
tables = _table_detector.detect(image)
|
||||
|
||||
if not tables:
|
||||
return tables
|
||||
|
||||
recognizer = TableStructureRecognizer()
|
||||
for table in tables:
|
||||
if table.image_crop is not None:
|
||||
structure, markdown = recognizer.recognize(table.image_crop)
|
||||
structure, markdown = _table_structure_recognizer.recognize(table.image_crop)
|
||||
table.structure = structure
|
||||
table.markdown = markdown
|
||||
|
||||
|
||||
@@ -63,6 +63,7 @@ class OCRResult:
|
||||
boxes: list[dict[str, Any]]
|
||||
raw_result: Any
|
||||
processing_time: float
|
||||
fallback_reason: str | None = None
|
||||
|
||||
|
||||
class PaddleOCREngine:
|
||||
@@ -78,17 +79,23 @@ class PaddleOCREngine:
|
||||
|
||||
paddle_lang = map_languages(settings.ocr.languages, "paddleocr")
|
||||
|
||||
self._ocr = PaddleOCR(
|
||||
use_angle_cls=settings.ocr.use_angle_cls,
|
||||
lang=paddle_lang,
|
||||
use_gpu=settings.ocr.use_gpu,
|
||||
show_log=False,
|
||||
det_db_thresh=settings.ocr.det_db_thresh,
|
||||
det_db_box_thresh=settings.ocr.det_db_box_thresh,
|
||||
det_db_unclip_ratio=settings.ocr.det_db_unclip_ratio,
|
||||
rec_batch_num=settings.ocr.rec_batch_num,
|
||||
cpu_threads=settings.ocr.cpu_threads,
|
||||
)
|
||||
legacy_options = {
|
||||
"use_angle_cls": settings.ocr.use_angle_cls,
|
||||
"lang": paddle_lang,
|
||||
"use_gpu": settings.ocr.use_gpu,
|
||||
"show_log": False,
|
||||
"det_db_thresh": settings.ocr.det_db_thresh,
|
||||
"det_db_box_thresh": settings.ocr.det_db_box_thresh,
|
||||
"det_db_unclip_ratio": settings.ocr.det_db_unclip_ratio,
|
||||
"rec_batch_num": settings.ocr.rec_batch_num,
|
||||
"cpu_threads": settings.ocr.cpu_threads,
|
||||
}
|
||||
try:
|
||||
self._ocr = PaddleOCR(**legacy_options)
|
||||
except TypeError as exc:
|
||||
# PaddleOCR 3.x removed several 2.x constructor parameters.
|
||||
logger.warning("paddleocr_legacy_api_rejected", error=str(exc))
|
||||
self._ocr = PaddleOCR(lang=paddle_lang)
|
||||
self._initialized = True
|
||||
logger.info("paddleocr_initialized", language=paddle_lang)
|
||||
except Exception as e:
|
||||
|
||||
@@ -26,7 +26,8 @@ class ProcessedImage:
|
||||
def _process_single_image(args: tuple[Path, dict]) -> ProcessedImage:
|
||||
"""Worker function for multiprocessing - must be at module level."""
|
||||
path, config_dict = args
|
||||
from ocr_pipeline.ocr.engine import ImagePreprocessor, OCREngine, compute_image_hash
|
||||
from ocr_pipeline.ocr.engine import OCREngine
|
||||
from ocr_pipeline.ocr.preprocess import ImagePreprocessor, compute_image_hash
|
||||
|
||||
preprocessor = ImagePreprocessor()
|
||||
ocr_engine = OCREngine()
|
||||
@@ -86,7 +87,13 @@ class ParallelProcessor:
|
||||
unique_images = []
|
||||
seen = set()
|
||||
for img in all_images:
|
||||
if img.is_file() and img.suffix.lower() 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
|
||||
# Path.match can vary for absolute paths; explicitly protect bundle directories.
|
||||
if any(part.endswith(".photoslibrary") for part in img.parts):
|
||||
continue
|
||||
if img not in seen:
|
||||
seen.add(img)
|
||||
unique_images.append(img)
|
||||
|
||||
@@ -1,119 +1,95 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from ocr_pipeline.config import settings
|
||||
from ocr_pipeline.utils.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class ImagePreprocessor:
|
||||
"""Image preprocessing for OCR optimization."""
|
||||
"""Creates an OCR-friendly derivative without mutating detector input images."""
|
||||
|
||||
def __init__(self):
|
||||
def __init__(self) -> None:
|
||||
self.config = settings.ocr.preprocess
|
||||
self.max_dim = self.config.max_dimension
|
||||
|
||||
def load_image(self, path: str) -> np.ndarray:
|
||||
"""Load image from file path."""
|
||||
def load_image(self, path: str | Path) -> np.ndarray:
|
||||
image = cv2.imread(str(path))
|
||||
if image is None:
|
||||
raise ValueError(f"Failed to load image: {path}")
|
||||
return image
|
||||
|
||||
def resize_if_needed(self, image: np.ndarray) -> np.ndarray:
|
||||
"""Resize image if it exceeds max dimension."""
|
||||
h, w = image.shape[:2]
|
||||
if max(h, w) > self.max_dim:
|
||||
scale = self.max_dim / max(h, w)
|
||||
new_w, new_h = int(w * scale), int(h * scale)
|
||||
image = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_AREA)
|
||||
height, width = image.shape[:2]
|
||||
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)
|
||||
|
||||
def deskew(self, image: np.ndarray) -> np.ndarray:
|
||||
"""Correct skew in the image."""
|
||||
if not self.config.deskew:
|
||||
return image
|
||||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||||
coords = np.column_stack(np.where(gray > 0))
|
||||
if len(coords) == 0:
|
||||
_, foreground = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
|
||||
coords = np.column_stack(np.where(foreground > 0))
|
||||
if len(coords) < 20:
|
||||
return image
|
||||
angle = cv2.minAreaRect(coords)[-1]
|
||||
if angle < -45:
|
||||
angle = -(90 + angle)
|
||||
else:
|
||||
angle = -angle
|
||||
if abs(angle) > 0.5:
|
||||
h, w = image.shape[:2]
|
||||
center = (w // 2, h // 2)
|
||||
M = cv2.getRotationMatrix2D(center, angle, 1.0)
|
||||
image = cv2.warpAffine(image, M, (w, h), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REPLICATE)
|
||||
angle = -(90 + angle) if angle < -45 else -angle
|
||||
if abs(angle) <= 0.5:
|
||||
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)
|
||||
|
||||
def denoise(self, image: np.ndarray) -> np.ndarray:
|
||||
"""Remove noise from image."""
|
||||
if not self.config.denoise:
|
||||
return image
|
||||
return cv2.fastNlMeansDenoisingColored(image, None, 10, 10, 7, 21)
|
||||
return cv2.fastNlMeansDenoisingColored(image, None, 10, 10, 7, 21) if self.config.denoise else image
|
||||
|
||||
def apply_clahe(self, image: np.ndarray) -> np.ndarray:
|
||||
"""Apply CLAHE for contrast enhancement."""
|
||||
if not self.config.clahe:
|
||||
return image
|
||||
lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
|
||||
l, a, b = cv2.split(lab)
|
||||
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
|
||||
l = clahe.apply(l)
|
||||
lab = cv2.merge((l, a, b))
|
||||
return cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
|
||||
lightness, a_channel, b_channel = cv2.split(lab)
|
||||
lightness = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)).apply(lightness)
|
||||
return cv2.cvtColor(cv2.merge((lightness, a_channel, b_channel)), cv2.COLOR_LAB2BGR)
|
||||
|
||||
def remove_lines(self, image: np.ndarray) -> np.ndarray:
|
||||
"""Remove horizontal and vertical lines (tables, grids)."""
|
||||
if not self.config.remove_lines:
|
||||
return image
|
||||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||||
_, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
|
||||
horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (40, 1))
|
||||
vertical_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 40))
|
||||
horizontal_lines = cv2.morphologyEx(binary, cv2.MORPH_OPEN, horizontal_kernel)
|
||||
vertical_lines = cv2.morphologyEx(binary, cv2.MORPH_OPEN, vertical_kernel)
|
||||
lines = cv2.bitwise_or(horizontal_lines, vertical_lines)
|
||||
contours, _ = cv2.findContours(lines, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
||||
mask = np.ones(image.shape[:2], dtype=np.uint8) * 255
|
||||
cv2.drawContours(mask, contours, -1, 0, thickness=cv2.FILLED)
|
||||
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]
|
||||
result[mask > 0] = (255, 255, 255)
|
||||
return result
|
||||
|
||||
def adaptive_threshold(self, image: np.ndarray) -> np.ndarray:
|
||||
"""Apply adaptive thresholding for binarization."""
|
||||
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(self, path: str) -> np.ndarray:
|
||||
"""Full preprocessing pipeline."""
|
||||
image = self.load_image(path)
|
||||
def preprocess_image(self, image: np.ndarray) -> np.ndarray:
|
||||
image = self.resize_if_needed(image)
|
||||
image = self.deskew(image)
|
||||
image = self.denoise(image)
|
||||
image = self.apply_clahe(image)
|
||||
image = self.remove_lines(image)
|
||||
image = self.adaptive_threshold(image)
|
||||
return image
|
||||
return self.adaptive_threshold(image)
|
||||
|
||||
def preprocess(self, path: str | Path) -> np.ndarray:
|
||||
return self.preprocess_image(self.load_image(path))
|
||||
|
||||
|
||||
def compute_image_hash(path: str) -> str:
|
||||
"""Compute SHA256 hash of image file for deduplication."""
|
||||
def compute_image_hash(path: str | Path) -> str:
|
||||
hasher = hashlib.sha256()
|
||||
with open(path, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(8192), b""):
|
||||
with Path(path).open("rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(8192), b""):
|
||||
hasher.update(chunk)
|
||||
return hasher.hexdigest()
|
||||
|
||||
@@ -3,13 +3,19 @@ from __future__ import annotations
|
||||
import datetime
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import yaml
|
||||
|
||||
from ocr_pipeline.config import settings
|
||||
from ocr_pipeline.utils.logging import get_logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ocr_pipeline.citations import Citation
|
||||
from ocr_pipeline.detectors.figures import DetectedFigure
|
||||
from ocr_pipeline.detectors.tables import DetectedTable
|
||||
from ocr_pipeline.postprocess.entities import EntityResult
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -30,131 +36,144 @@ class OutputMetadata:
|
||||
|
||||
|
||||
class MarkdownWriter:
|
||||
def __init__(self):
|
||||
def __init__(self) -> None:
|
||||
self.base_dir = Path(settings.output.base_directory).expanduser().resolve()
|
||||
self.organize_by = settings.output.organize_by
|
||||
self.frontmatter_fields = settings.output.frontmatter.fields
|
||||
self._run_dir: Path | None = None
|
||||
self._written_chunks: list[tuple[dict[str, Any], str]] = []
|
||||
|
||||
def _get_output_dir(self, timestamp: str) -> Path:
|
||||
def start_run(self, timestamp: str, run_number: int) -> Path:
|
||||
"""Allocate one shared output directory for the entire pipeline run."""
|
||||
self._written_chunks = []
|
||||
dt = datetime.datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
|
||||
|
||||
if self.organize_by == "date_run":
|
||||
date_dir = self.base_dir / dt.strftime("%Y-%m-%d")
|
||||
run_dirs = sorted(date_dir.glob("run_*"))
|
||||
run_num = len(run_dirs) + 1
|
||||
return date_dir / f"run_{run_num:03d}"
|
||||
elif self.organize_by == "source_dir":
|
||||
return self.base_dir / "by_source"
|
||||
self._run_dir = self.base_dir / dt.strftime("%Y-%m-%d") / f"run_{run_number:03d}"
|
||||
elif self.organize_by == "flat":
|
||||
self._run_dir = self.base_dir
|
||||
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, "/"))
|
||||
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:
|
||||
return self._run_dir
|
||||
timestamp = metadata.get("timestamp")
|
||||
if not isinstance(timestamp, str) or not timestamp:
|
||||
timestamp = datetime.datetime.now(datetime.UTC).isoformat()
|
||||
return self.start_run(timestamp, 1)
|
||||
|
||||
def write_chunk(
|
||||
self,
|
||||
*,
|
||||
content: str,
|
||||
metadata: dict[str, Any],
|
||||
figures: list[Any] = None,
|
||||
tables: list[Any] = None,
|
||||
entities: Any = None,
|
||||
citations: list[Any] = None,
|
||||
figures: list[DetectedFigure] | None = None,
|
||||
tables: list[DetectedTable] | None = None,
|
||||
entities: EntityResult | None = None,
|
||||
citations: list[Citation] | None = None,
|
||||
) -> Path:
|
||||
output_dir = self._get_output_dir(metadata.get("timestamp", ""))
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
source_path = Path(metadata.get("source_path", "unknown"))
|
||||
safe_name = "".join(c for c in source_path.stem if c.isalnum() or c in "-_")[:80]
|
||||
chunk_idx = metadata.get("chunk_index", 0)
|
||||
total_chunks = metadata.get("total_chunks", 1)
|
||||
|
||||
if total_chunks > 1:
|
||||
filename = f"{safe_name}_chunk_{chunk_idx:03d}.md"
|
||||
else:
|
||||
filename = f"{safe_name}.md"
|
||||
|
||||
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"
|
||||
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"
|
||||
output_path = output_dir / filename
|
||||
|
||||
frontmatter = self._build_frontmatter(metadata)
|
||||
body = self._build_body(content, figures, tables, entities, citations)
|
||||
|
||||
with open(output_path, "w") as f:
|
||||
f.write("---\n")
|
||||
yaml.dump(frontmatter, f, sort_keys=False, allow_unicode=True)
|
||||
f.write("---\n\n")
|
||||
f.write(body)
|
||||
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")
|
||||
yaml.safe_dump(frontmatter, handle, sort_keys=False, allow_unicode=True)
|
||||
handle.write("---\n\n")
|
||||
handle.write(body)
|
||||
|
||||
logger.debug("markdown_written", path=str(output_path))
|
||||
return output_path
|
||||
|
||||
def write_consolidated(self) -> Path | None:
|
||||
"""Write one readable document while retaining chunk files for RAG."""
|
||||
if not self._written_chunks:
|
||||
return None
|
||||
output_dir = self._run_dir or self.base_dir
|
||||
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))
|
||||
|
||||
timestamp = datetime.datetime.now(datetime.UTC).isoformat()
|
||||
frontmatter = {
|
||||
"document_type": "consolidated_ocr",
|
||||
"generated_at": timestamp,
|
||||
"source_count": len(grouped),
|
||||
"chunk_count": len(self._written_chunks),
|
||||
"sources": [str(path) for path in grouped],
|
||||
}
|
||||
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))):
|
||||
parts.extend([content.strip(), ""])
|
||||
output_path = output_dir / settings.output.consolidated_filename
|
||||
with output_path.open("w", encoding="utf-8") as handle:
|
||||
handle.write("---\n")
|
||||
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))
|
||||
return output_path
|
||||
|
||||
def _build_frontmatter(self, metadata: dict[str, Any]) -> dict[str, Any]:
|
||||
fm = {}
|
||||
for field in self.frontmatter_fields:
|
||||
if field in metadata:
|
||||
fm[field] = metadata[field]
|
||||
return fm
|
||||
return {field: metadata[field] for field in self.frontmatter_fields if field in metadata}
|
||||
|
||||
def _build_body(
|
||||
self,
|
||||
content: str,
|
||||
figures: list[Any] = None,
|
||||
tables: list[Any] = None,
|
||||
entities: Any = None,
|
||||
citations: list[Any] = None,
|
||||
source_name: str,
|
||||
figures: list[DetectedFigure],
|
||||
tables: list[DetectedTable],
|
||||
entities: EntityResult | None,
|
||||
citations: list[Citation],
|
||||
) -> str:
|
||||
parts = []
|
||||
|
||||
source_name = Path(content.split("\n")[0]).name if content else "Unknown"
|
||||
parts.append(f"# Screenshot: {source_name}")
|
||||
parts.append("")
|
||||
|
||||
if content:
|
||||
parts.append("## Extracted Text")
|
||||
parts.append("")
|
||||
parts.append(content.strip())
|
||||
parts.append("")
|
||||
|
||||
parts = [f"# Screenshot: {source_name}", "", "## Extracted Text", "", content.strip(), ""]
|
||||
if figures:
|
||||
parts.append("## Figures Detected")
|
||||
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}"])
|
||||
if figure.caption:
|
||||
parts.append(f"- **Caption**: {figure.caption}")
|
||||
parts.append("")
|
||||
for i, fig in enumerate(figures):
|
||||
parts.append(f"### Figure {i + 1}")
|
||||
parts.append(f"- **Bounding Box**: {fig.bbox}")
|
||||
parts.append(f"- **Confidence**: {fig.confidence:.2f}")
|
||||
if fig.caption:
|
||||
parts.append(f"- **Caption**: {fig.caption}")
|
||||
parts.append("")
|
||||
|
||||
if tables:
|
||||
parts.append("## Tables Detected")
|
||||
parts.append("")
|
||||
for i, table in enumerate(tables):
|
||||
parts.append(f"### Table {i + 1}")
|
||||
parts.append(f"- **Bounding Box**: {table.bbox}")
|
||||
parts.append(f"- **Confidence**: {table.confidence:.2f}")
|
||||
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}"])
|
||||
if table.markdown:
|
||||
parts.extend(["", "#### Table Content (Markdown)", "", table.markdown])
|
||||
parts.append("")
|
||||
parts.append("#### Table Content (Markdown)")
|
||||
parts.append("")
|
||||
parts.append(table.markdown)
|
||||
parts.append("")
|
||||
|
||||
if entities and entities.entities:
|
||||
parts.append("## Detected Entities")
|
||||
parts.extend(["## Detected Entities", "", "| Entity | Type | Context |", "|--------|------|---------|"])
|
||||
for entity in entities.entities[:50]:
|
||||
context = entity.context[:80].replace("|", "\\|")
|
||||
if len(entity.context) > 80:
|
||||
context += "..."
|
||||
parts.append(f"| {entity.text} | {entity.label} | {context} |")
|
||||
parts.append("")
|
||||
parts.append("| Entity | Type | Context |")
|
||||
parts.append("|--------|------|---------|")
|
||||
for ent in entities.entities[:50]:
|
||||
context = ent.context[:80].replace("|", "\\|") + (
|
||||
"..." if len(ent.context) > 80 else ""
|
||||
)
|
||||
parts.append(f"| {ent.text} | {ent.label} | {context} |")
|
||||
parts.append("")
|
||||
|
||||
if citations:
|
||||
parts.append("## Citations Found")
|
||||
parts.extend(["## Citations Found", "", "| Type | Identifier |", "|------|------------|"])
|
||||
for citation in citations:
|
||||
parts.append(f"| {citation.type.upper()} | `{citation.identifier}` |")
|
||||
parts.append("")
|
||||
parts.append("| Type | Identifier |")
|
||||
parts.append("|------|------------|")
|
||||
for cit in citations:
|
||||
parts.append(f"| {cit.type.upper()} | `{cit.identifier}` |")
|
||||
parts.append("")
|
||||
|
||||
return "\n".join(parts)
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
import time
|
||||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
@@ -13,17 +13,10 @@ from ocr_pipeline.ocr import ImagePreprocessor, OCREngine, compute_image_hash
|
||||
from ocr_pipeline.ocr.parallel import ParallelProcessor
|
||||
from ocr_pipeline.output import MarkdownWriter
|
||||
from ocr_pipeline.postprocess import chunk_text, clean_text, extract_entities
|
||||
from ocr_pipeline.utils.db import get_db
|
||||
from ocr_pipeline.utils.db import Database, get_db
|
||||
from ocr_pipeline.utils.logging import get_logger
|
||||
from ocr_pipeline.watch import Watcher
|
||||
|
||||
try:
|
||||
from rich.console import Console
|
||||
|
||||
_console = Console()
|
||||
except ImportError:
|
||||
_console = None
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -32,157 +25,104 @@ class PipelineResult:
|
||||
total_images: int = 0
|
||||
successful: int = 0
|
||||
failed: int = 0
|
||||
skipped: int = 0
|
||||
chunks_created: int = 0
|
||||
output_files: list[Path] = field(default_factory=list)
|
||||
errors: list[dict[str, Any]] = field(default_factory=list)
|
||||
errors: list[dict[str, str]] = field(default_factory=list)
|
||||
|
||||
|
||||
class OCRPipeline:
|
||||
def __init__(self):
|
||||
def __init__(self) -> None:
|
||||
self.engine = OCREngine()
|
||||
self.preprocessor = ImagePreprocessor()
|
||||
self.writer = MarkdownWriter()
|
||||
self.db = get_db()
|
||||
|
||||
def _process_single(self, image_path: Path) -> list[dict[str, Any]]:
|
||||
image = self.preprocessor.preprocess(image_path)
|
||||
result = self.engine.process(image)
|
||||
self.db: Database = get_db()
|
||||
self._run_id: int | None = None
|
||||
|
||||
def _process_single(self, image_path: Path) -> list[Path]:
|
||||
image_hash = compute_image_hash(image_path)
|
||||
if self.db.is_processed(image_hash, image_path):
|
||||
logger.debug("already_processed", path=str(image_path))
|
||||
if settings.processing.skip_existing and self.db.is_processed(image_hash, image_path):
|
||||
logger.info("already_processed", path=str(image_path))
|
||||
return []
|
||||
|
||||
cleaned = clean_text(result.text)
|
||||
original_image = self.preprocessor.load_image(image_path)
|
||||
ocr_image = self.preprocessor.preprocess_image(original_image)
|
||||
ocr_result = self.engine.process(ocr_image)
|
||||
cleaned = clean_text(ocr_result.text)
|
||||
chunks = chunk_text(cleaned.text)
|
||||
entities = extract_entities(cleaned.text)
|
||||
figures = detect_figures(image)
|
||||
tables = detect_tables(image)
|
||||
# Detectors receive the original image: binarization/line removal damages their inputs.
|
||||
figures = detect_figures(original_image)
|
||||
tables = detect_tables(original_image)
|
||||
citations = extract_citations(cleaned.text)
|
||||
timestamp = datetime.datetime.now(datetime.UTC).isoformat()
|
||||
|
||||
output_files = []
|
||||
for i, chunk in enumerate(chunks):
|
||||
output_path = self.writer.write_chunk(
|
||||
source_path=image_path,
|
||||
source_hash=image_hash,
|
||||
text=chunk.content,
|
||||
entities=[e.text for e in entities.entities],
|
||||
figures=figures,
|
||||
tables=tables,
|
||||
citations=[f"{c.type}:{c.identifier}" for c in citations],
|
||||
chunk_index=i,
|
||||
total_chunks=len(chunks),
|
||||
ocr_engine=result.engine,
|
||||
ocr_confidence=result.confidence,
|
||||
language=result.language,
|
||||
)
|
||||
output_files.append(output_path)
|
||||
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),
|
||||
}
|
||||
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=result.engine,
|
||||
confidence=result.confidence,
|
||||
chunks=len(chunks),
|
||||
run_id=0,
|
||||
)
|
||||
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 [{"path": p, "chunks": len(chunks)} for p in output_files]
|
||||
def process_single(self, image_path: Path) -> list[Path]:
|
||||
"""Public single-image API used by notebook migration and watch mode."""
|
||||
return self._process_single(image_path)
|
||||
|
||||
def process_all(self, image_paths: list[Path] | None = None) -> PipelineResult:
|
||||
"""Process images through the full pipeline."""
|
||||
if image_paths is None:
|
||||
processor = ParallelProcessor()
|
||||
results = processor.process()
|
||||
image_paths = [r.path for r in results if r.success]
|
||||
image_paths = ParallelProcessor().find_images()
|
||||
else:
|
||||
# Validate paths exist
|
||||
image_paths = [p for p in image_paths if p.exists()]
|
||||
|
||||
image_paths = [path for path in image_paths if path.exists()]
|
||||
result = PipelineResult(total_images=len(image_paths))
|
||||
|
||||
if not image_paths:
|
||||
return result
|
||||
|
||||
with ProcessPoolExecutor(max_workers=settings.processing.workers) as executor:
|
||||
future_to_path = {executor.submit(self._process_single, p): p for p in image_paths}
|
||||
|
||||
for future in as_completed(future_to_path):
|
||||
path = future_to_path[future]
|
||||
run_timestamp = datetime.datetime.now(datetime.UTC).isoformat()
|
||||
self._run_id, run_number = self.db.start_run()
|
||||
self.writer.start_run(run_timestamp, run_number)
|
||||
try:
|
||||
output_info = future.result(timeout=settings.processing.timeout_per_image)
|
||||
# Output/database writes are deliberately centralized to avoid collisions and SQLite contention.
|
||||
for path in image_paths:
|
||||
try:
|
||||
outputs = self._process_single(path)
|
||||
if not outputs and settings.processing.skip_existing:
|
||||
result.skipped += 1
|
||||
else:
|
||||
result.successful += 1
|
||||
result.chunks_created += sum(info["chunks"] for info in output_info)
|
||||
result.output_files.extend([info["path"] for info in output_info])
|
||||
except Exception as e:
|
||||
result.chunks_created += len(outputs)
|
||||
result.output_files.extend(outputs)
|
||||
except Exception as exc:
|
||||
result.failed += 1
|
||||
result.errors.append({"path": str(path), "error": str(e)})
|
||||
logger.error("pipeline_failed", path=str(path), error=str(e))
|
||||
|
||||
result.errors.append({"path": str(path), "error": str(exc)})
|
||||
logger.exception("pipeline_failed", path=str(path))
|
||||
if settings.output.write_consolidated:
|
||||
consolidated = self.writer.write_consolidated()
|
||||
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._run_id = None
|
||||
return result
|
||||
|
||||
def watch(self):
|
||||
def process_callback(path: Path):
|
||||
self._process_single(path)
|
||||
def watch(self) -> None:
|
||||
def process_callback(path: Path) -> None:
|
||||
self.process_single(path)
|
||||
|
||||
watcher = Watcher(process_callback)
|
||||
watcher.start()
|
||||
|
||||
try:
|
||||
while True:
|
||||
time.sleep(1)
|
||||
except KeyboardInterrupt:
|
||||
watcher.stop()
|
||||
|
||||
def get_stats(self) -> dict[str, Any]:
|
||||
db_stats = self.db.get_stats()
|
||||
return {
|
||||
"total_files_processed": db_stats.get("total_processed", 0),
|
||||
"total_runs": db_stats.get("total_runs", 0),
|
||||
"latest_run": db_stats.get("latest_run"),
|
||||
}
|
||||
|
||||
|
||||
def migrate_notebook(notebook_path: Path, pipeline: OCRPipeline) -> dict[str, Any]:
|
||||
import json
|
||||
|
||||
with open(notebook_path) as f:
|
||||
nb = json.load(f)
|
||||
|
||||
image_paths = set()
|
||||
for cell in nb.get("cells", []):
|
||||
if cell.get("cell_type") == "code":
|
||||
for output in cell.get("outputs", []):
|
||||
if output.get("output_type") == "stream":
|
||||
for line in output.get("text", []):
|
||||
parts = line.split()
|
||||
for p in parts:
|
||||
if "SCR-" in p and (
|
||||
p.endswith(".png") or p.endswith(".jpg") or p.endswith(".jpeg")
|
||||
):
|
||||
image_paths.add(p)
|
||||
|
||||
unique_paths = [Path(p) for p in image_paths if Path(p).exists()]
|
||||
|
||||
try:
|
||||
from rich.console import Console
|
||||
|
||||
console = Console()
|
||||
console.print(f"Found {len(unique_paths)} images in notebook")
|
||||
except ImportError:
|
||||
logger.info("Found %d images in notebook", len(unique_paths))
|
||||
|
||||
if not unique_paths:
|
||||
return {"total": 0, "successful": 0, "failed": 0, "errors": []}
|
||||
|
||||
result = pipeline.process_all(unique_paths)
|
||||
|
||||
return {
|
||||
"total": result.total_images,
|
||||
"successful": result.successful,
|
||||
"failed": result.failed,
|
||||
"chunks_created": result.chunks_created,
|
||||
"output_files": len(result.output_files),
|
||||
"errors": result.errors,
|
||||
}
|
||||
def get_stats(self, days: int = 30) -> dict[str, Any]:
|
||||
return self.db.get_stats(days)
|
||||
|
||||
@@ -2,7 +2,7 @@ from __future__ import annotations
|
||||
|
||||
from .chunk import SemanticChunker, TextChunk, chunk_text
|
||||
from .clean import CleanResult, clean_text
|
||||
from .entities import Entity, SciSpacyNER, extract_entities, summarize_entities
|
||||
from .entities import Entity, ScientificEntityRecognizer, extract_entities, summarize_entities
|
||||
|
||||
__all__ = [
|
||||
"clean_text",
|
||||
@@ -13,5 +13,5 @@ __all__ = [
|
||||
"extract_entities",
|
||||
"summarize_entities",
|
||||
"Entity",
|
||||
"SciSpacyNER",
|
||||
"ScientificEntityRecognizer",
|
||||
]
|
||||
|
||||
@@ -3,7 +3,7 @@ from __future__ import annotations
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||||
|
||||
from ocr_pipeline.config import settings
|
||||
from ocr_pipeline.utils.logging import get_logger
|
||||
|
||||
@@ -24,140 +24,54 @@ def fix_encoding(text: str) -> str:
|
||||
return text
|
||||
|
||||
|
||||
def normalize_unicode(text: str) -> str:
|
||||
return unicodedata.normalize("NFC", text)
|
||||
|
||||
|
||||
def fix_ligatures(text: str) -> str:
|
||||
ligatures = {
|
||||
"fi": "fi",
|
||||
"fl": "fl",
|
||||
"ff": "ff",
|
||||
"ffi": "ffi",
|
||||
"ffl": "ffl",
|
||||
"ſt": "ft",
|
||||
"st": "st",
|
||||
"Ꜳ": "AA",
|
||||
"ꜳ": "aa",
|
||||
"Ꜵ": "AO",
|
||||
"ꜵ": "ao",
|
||||
"Ꜷ": "AU",
|
||||
"ꜷ": "au",
|
||||
"Ꜹ": "AV",
|
||||
"ꜹ": "av",
|
||||
"Ꜻ": "AV",
|
||||
"ꜻ": "av",
|
||||
"Ꜽ": "AY",
|
||||
"ꜽ": "ay",
|
||||
}
|
||||
for lig, repl in ligatures.items():
|
||||
text = text.replace(lig, repl)
|
||||
return text
|
||||
|
||||
|
||||
def remove_control_chars(text: str) -> str:
|
||||
return "".join(ch for ch in text if unicodedata.category(ch)[0] != "C" or ch in "\n\t")
|
||||
|
||||
|
||||
def fix_hyphenation(text: str) -> str:
|
||||
return re.sub(r"(\w+)-\n(\w+)", r"\1\2", text)
|
||||
|
||||
|
||||
def normalize_whitespace(text: str) -> str:
|
||||
text = re.sub(r"[ \t]+", " ", text)
|
||||
text = re.sub(r"\n{3,}", "\n\n", text)
|
||||
return text.strip()
|
||||
|
||||
|
||||
def remove_ocr_artifacts(text: str) -> str:
|
||||
patterns = [
|
||||
r"^\s*[\|┃║]\s*$",
|
||||
r"^\s*[─━]\s*$",
|
||||
r"^\s*[┌┐└┘├┤┬┴┼]\s*$",
|
||||
r"(.)\1{10,}",
|
||||
r"^\s*[•·▪▫◦‣⁃]\s*$",
|
||||
]
|
||||
lines = text.split("\n")
|
||||
cleaned = []
|
||||
for line in lines:
|
||||
if any(re.match(p, line.strip()) for p in patterns):
|
||||
continue
|
||||
cleaned.append(line)
|
||||
return "\n".join(cleaned)
|
||||
|
||||
|
||||
def fix_scientific_notation(text: str) -> str:
|
||||
text = re.sub(r"(\d)\s*[×xX]\s*10\s*[\^]?\s*([+-]?\d+)", r"\1×10^\2", text)
|
||||
text = re.sub(r"(\d)\s*[×xX]\s*10\s*([+-]?\d+)", r"\1×10^\2", text)
|
||||
return text
|
||||
|
||||
|
||||
def normalize_units(text: str) -> str:
|
||||
unit_fixes = {
|
||||
r"(\d)\s*u[mM]\b": r"\1 µm",
|
||||
r"(\d)\s*[µu]g\b": r"\1 µg",
|
||||
r"(\d)\s*[mn]g\b": r"\1 mg",
|
||||
r"(\d)\s*[kK]b\b": r"\1 kb",
|
||||
r"(\d)\s*[mM]b\b": r"\1 Mb",
|
||||
r"(\d)\s*[gG]b\b": r"\1 Gb",
|
||||
r"(\d)\s*[pP]?[mM]\b": r"\1 pm",
|
||||
r"(\d)\s*[nN][mM]\b": r"\1 nm",
|
||||
r"(\d)\s*[cC][mM]\b": r"\1 cm",
|
||||
r"(\d)\s*[mM][lL]\b": r"\1 mL",
|
||||
r"(\d)\s*[uU][lL]\b": r"\1 µL",
|
||||
r"(\d)\s*°\s*[cC]\b": r"\1 °C",
|
||||
}
|
||||
for pattern, repl in unit_fixes.items():
|
||||
text = re.sub(pattern, repl, text)
|
||||
return text
|
||||
|
||||
|
||||
def fix_line_breaks(text: str) -> str:
|
||||
text = re.sub(r"(?<!\n)\n(?!\n)", " ", text)
|
||||
text = re.sub(r"\n{3,}", "\n\n", text)
|
||||
number = r"(\d+(?:\.\d+)?)"
|
||||
replacements = [
|
||||
(rf"{number}\s*[uµ]M\b", r"\1 µM"),
|
||||
(rf"{number}\s*[uµ]m\b", r"\1 µm"),
|
||||
(rf"{number}\s*[uµ]g\b", r"\1 µg"),
|
||||
(rf"{number}\s*ng\b", r"\1 ng"),
|
||||
(rf"{number}\s*mg\b", r"\1 mg"),
|
||||
(rf"{number}\s*[uµ]L\b", r"\1 µL"),
|
||||
(rf"{number}\s*mL\b", r"\1 mL"),
|
||||
(rf"{number}\s*pM\b", r"\1 pM"),
|
||||
(rf"{number}\s*nM\b", r"\1 nM"),
|
||||
(rf"{number}\s*mM\b", r"\1 mM"),
|
||||
(rf"{number}\s*[kK]b\b", r"\1 kb"),
|
||||
(rf"{number}\s*[mM]b\b", r"\1 Mb"),
|
||||
(rf"{number}\s*[gG]b\b", r"\1 Gb"),
|
||||
(rf"{number}\s*°\s*[cC]\b", r"\1 °C"),
|
||||
]
|
||||
for pattern, replacement in replacements:
|
||||
text = re.sub(pattern, replacement, text)
|
||||
return text
|
||||
|
||||
|
||||
def clean_text(text: str) -> CleanResult:
|
||||
original_len = len(text)
|
||||
operations = []
|
||||
|
||||
original_length = len(text)
|
||||
operations: list[str] = []
|
||||
try:
|
||||
text = fix_encoding(text)
|
||||
finally:
|
||||
operations.append("fix_encoding")
|
||||
|
||||
text = normalize_unicode(text)
|
||||
text = unicodedata.normalize("NFC", text)
|
||||
operations.append("normalize_unicode")
|
||||
|
||||
text = fix_ligatures(text)
|
||||
for source, target in {"fi": "fi", "fl": "fl", "ff": "ff", "ffi": "ffi", "ffl": "ffl"}.items():
|
||||
text = text.replace(source, target)
|
||||
operations.append("fix_ligatures")
|
||||
|
||||
text = remove_control_chars(text)
|
||||
text = "".join(char for char in text if unicodedata.category(char)[0] != "C" or char in "\n\t")
|
||||
operations.append("remove_control_chars")
|
||||
|
||||
text = fix_hyphenation(text)
|
||||
text = re.sub(r"(\w+)-\n(\w+)", r"\1\2", text)
|
||||
operations.append("fix_hyphenation")
|
||||
|
||||
text = normalize_whitespace(text)
|
||||
operations.append("normalize_whitespace")
|
||||
|
||||
text = remove_ocr_artifacts(text)
|
||||
text = re.sub(r"(?m)^\s*[|┃║─━┌┐└┘├┤┬┴┼]+\s*$", "", text)
|
||||
text = re.sub(r"(.)\1{10,}", "", text)
|
||||
operations.append("remove_ocr_artifacts")
|
||||
|
||||
text = fix_scientific_notation(text)
|
||||
text = re.sub(r"(\d)\s*[×xX]\s*10\s*\^?\s*([+-]?\d+)", r"\1×10^\2", text)
|
||||
operations.append("fix_scientific_notation")
|
||||
|
||||
text = normalize_units(text)
|
||||
operations.append("normalize_units")
|
||||
|
||||
text = fix_line_breaks(text)
|
||||
operations.append("fix_line_breaks")
|
||||
|
||||
cleaned_len = len(text)
|
||||
logger.debug("text_cleaned", original=original_len, cleaned=cleaned_len, ops=operations)
|
||||
|
||||
return CleanResult(
|
||||
text=text,
|
||||
original_length=original_len,
|
||||
cleaned_length=cleaned_len,
|
||||
operations=operations,
|
||||
)
|
||||
text = re.sub(r"[ \t]+", " ", text)
|
||||
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)
|
||||
|
||||
@@ -27,12 +27,14 @@ class EntityResult:
|
||||
text: str
|
||||
total_entities: int
|
||||
entities_by_type: dict[str, int]
|
||||
backend: str = "unknown"
|
||||
|
||||
|
||||
class ScientificEntityRecognizer:
|
||||
def __init__(self):
|
||||
self._nlp = None
|
||||
self._initialized = False
|
||||
self.backend = "unavailable"
|
||||
|
||||
def _init(self):
|
||||
if self._initialized:
|
||||
@@ -42,6 +44,7 @@ class ScientificEntityRecognizer:
|
||||
self._nlp = spacy.load(settings.entities.model)
|
||||
self._nlp.max_length = 2_000_000
|
||||
self._initialized = True
|
||||
self.backend = "scispacy"
|
||||
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)
|
||||
@@ -52,6 +55,7 @@ class ScientificEntityRecognizer:
|
||||
try:
|
||||
self._nlp = spacy.blank("en")
|
||||
self._nlp.add_pipe("sentencizer")
|
||||
self.backend = "blank_spacy"
|
||||
logger.info("using_fallback_blank_model")
|
||||
except Exception:
|
||||
self._nlp = None
|
||||
@@ -99,6 +103,7 @@ class ScientificEntityRecognizer:
|
||||
text=text,
|
||||
total_entities=len(entities),
|
||||
entities_by_type=by_type,
|
||||
backend=self.backend,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("entity_extraction_failed", error=str(e))
|
||||
@@ -141,15 +146,14 @@ class ScientificEntityRecognizer:
|
||||
class RegexEntityRecognizer:
|
||||
PATTERNS = {
|
||||
"GENE": [
|
||||
r"\b(?:[A-Z]{2,5}\d*[A-Z]?|[A-Z]{3,10})\b",
|
||||
r"\b(?:BRCA[12]|TP53|EGFR|KRAS|MYC|PIK3CA|PTEN|APC|MLH1|MSH2|MSH6|PMS2|ATM|CHEK2|PALB2|RAD51|BARD1|BRIP1|CDH1|STK11|VHL|WT1|NF1|NF2|RB1|MEN1|RET|MET|ALK|ROS1|NTRK[123]|BRAF|NRAS|HRAS|IDH[12]|TERT|FGFR[1234]|KIT|PDGFRA|FLT3|JAK2|MPL|CALR|ASXL1|DNMT3A|TET2|SRSF2|U2AF1|SF3B1|ZRSR2|CBL|CSF3R|RUNX1|CEBPA|NPM1|FLT3|DNMT3A|IDH[12]|TET2|ASXL1|SRSF2|U2AF1|SF3B1|ZRSR2|CBL|CSF3R|RUNX1|CEBPA|NPM1)\b",
|
||||
r"\b(?:BRCA[12]|TP53|EGFR|KRAS|MYC|PIK3CA|PTEN|APC|MLH1|MSH2|MSH6|PMS2|ATM|CHEK2|PALB2|RAD51|BARD1|BRIP1|CDH1|STK11|VHL|WT1|NF1|NF2|RB1|MEN1|RET|MET|ALK|ROS1|NTRK[123]|BRAF|NRAS|HRAS|IDH[12]|TERT|FGFR[1234]|KIT|PDGFRA|FLT3|JAK2|MPL|CALR|ASXL1|DNMT3A|TET2|SRSF2|U2AF1|SF3B1|ZRSR2|CBL|CSF3R|RUNX1|CEBPA|NPM1)\b",
|
||||
],
|
||||
"PROTEIN": [
|
||||
r"\b(?:p53|p21|p16|CDK[1246]|cyclin\s+[ADEB]|Rb|E2F|MDM2|BAX|BCL[2XL]|CASP[389]|PARP|H2AX|ATM|ATR|CHK[12]|BRCA[12]|RAD51|FANCD2|FANCI|PCNA|RPA|MSH[26]|MLH1|PMS2|EXO1|MRE11|RAD50|NBN|CTIP|BRCA1|BRCA2|PALB2|RAD51C|RAD51D|BRIP1|BARD1)\b",
|
||||
],
|
||||
"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(?:[A-Z][a-z]?\d*(?:\([^)]+\))?(?:\s*[+\-]\s*[A-Z][a-z]?\d*(?:\([^)]+\))?)*)\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",
|
||||
@@ -177,9 +181,14 @@ class RegexEntityRecognizer:
|
||||
return entities
|
||||
|
||||
|
||||
_scientific_recognizer = ScientificEntityRecognizer()
|
||||
|
||||
|
||||
def extract_entities(text: str) -> EntityResult:
|
||||
recognizer = ScientificEntityRecognizer()
|
||||
result = recognizer.extract(text)
|
||||
if not settings.entities.enabled:
|
||||
return EntityResult(entities=[], text=text, total_entities=0, entities_by_type={})
|
||||
|
||||
result = _scientific_recognizer.extract(text)
|
||||
|
||||
regex_entities = RegexEntityRecognizer.extract(text)
|
||||
all_entities = result.entities + regex_entities
|
||||
@@ -194,6 +203,7 @@ def extract_entities(text: str) -> EntityResult:
|
||||
text=text,
|
||||
total_entities=len(all_entities),
|
||||
entities_by_type=by_type,
|
||||
backend=f"{result.backend}+regex" if result.backend != "unavailable" else "regex",
|
||||
)
|
||||
|
||||
|
||||
@@ -225,3 +235,13 @@ def format_entities_for_markdown(result: EntityResult) -> str:
|
||||
lines.append(f"| {ents[0].text} | {label} | {len(ents)} | {'; '.join(contexts)} |")
|
||||
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def summarize_entities(result: EntityResult) -> str:
|
||||
if not result.entities:
|
||||
return "No entities detected."
|
||||
|
||||
parts = [f"Total entities: {result.total_entities}"]
|
||||
for label, count in sorted(result.entities_by_type.items()):
|
||||
parts.append(f" {label}: {count}")
|
||||
return "\n".join(parts)
|
||||
|
||||
@@ -2,16 +2,14 @@ from __future__ import annotations
|
||||
|
||||
from .db import Database, ProcessedFile, get_db
|
||||
from .logging import ProgressTracker, create_progress, get_logger, setup_logging
|
||||
from .migrate import compute_file_hash, migrate_notebook, parse_notebook_for_images
|
||||
|
||||
|
||||
__all__ = [
|
||||
"setup_logging",
|
||||
"get_logger",
|
||||
"create_progress",
|
||||
"ProgressTracker",
|
||||
"parse_notebook_for_images",
|
||||
"migrate_notebook",
|
||||
"compute_file_hash",
|
||||
|
||||
"get_db",
|
||||
"Database",
|
||||
"ProcessedFile",
|
||||
|
||||
@@ -2,6 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import sqlite3
|
||||
import time
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
@@ -12,32 +13,34 @@ from ocr_pipeline.utils.logging import get_logger
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
@dataclass(frozen=True)
|
||||
class ProcessedFile:
|
||||
"""Legacy value object retained for import compatibility."""
|
||||
path: str
|
||||
hash: str
|
||||
timestamp: float
|
||||
run_id: int
|
||||
size: int
|
||||
mtime: float
|
||||
run_id: int | None = None
|
||||
size: int = 0
|
||||
mtime: float = 0.0
|
||||
|
||||
|
||||
class Database:
|
||||
def __init__(self, db_path: Path):
|
||||
self.db_path = Path(db_path).expanduser().resolve()
|
||||
def __init__(self, db_path: Path) -> None:
|
||||
self.db_path = db_path.expanduser().resolve()
|
||||
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._init_db()
|
||||
|
||||
@contextmanager
|
||||
def _conn(self):
|
||||
def _conn(self) -> Iterator[sqlite3.Connection]:
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
conn.row_factory = sqlite3.Row
|
||||
try:
|
||||
yield conn
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def _init_db(self):
|
||||
def _init_db(self) -> None:
|
||||
with self._conn() as conn:
|
||||
conn.executescript("""
|
||||
CREATE TABLE IF NOT EXISTS processed_files (
|
||||
@@ -47,9 +50,9 @@ class Database:
|
||||
ocr_engine TEXT,
|
||||
confidence REAL,
|
||||
chunks_created INTEGER DEFAULT 0,
|
||||
run_id INTEGER,
|
||||
created_at REAL DEFAULT (strftime('%s', 'now'))
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS runs (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
date TEXT NOT NULL,
|
||||
@@ -62,96 +65,49 @@ class Database:
|
||||
completed_at REAL,
|
||||
UNIQUE(date, run_number)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_processed_hash ON processed_files(hash);
|
||||
CREATE INDEX IF NOT EXISTS idx_runs_date ON runs(date);
|
||||
""")
|
||||
columns = {row[1] for row in conn.execute("PRAGMA table_info(processed_files)")}
|
||||
if "run_id" not in columns:
|
||||
conn.execute("ALTER TABLE processed_files ADD COLUMN run_id INTEGER")
|
||||
|
||||
def is_processed(self, file_hash: str, path: Path | str | None = None) -> bool:
|
||||
"""Check if a file is already processed by hash (and optionally path)."""
|
||||
with self._conn() as conn:
|
||||
if path:
|
||||
cursor = conn.execute(
|
||||
"SELECT 1 FROM processed_files WHERE hash = ? AND path = ?",
|
||||
(file_hash, str(path)),
|
||||
)
|
||||
if path is None:
|
||||
row = conn.execute("SELECT 1 FROM processed_files WHERE hash = ?", (file_hash,)).fetchone()
|
||||
else:
|
||||
cursor = conn.execute(
|
||||
"SELECT 1 FROM processed_files WHERE hash = ?",
|
||||
(file_hash,),
|
||||
)
|
||||
return cursor.fetchone() is not None
|
||||
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:
|
||||
cursor = conn.execute("SELECT * FROM processed_files WHERE hash = ?", (file_hash,))
|
||||
row = cursor.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 = 0,
|
||||
):
|
||||
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
|
||||
(hash, path, timestamp, ocr_engine, confidence, chunks_created)
|
||||
VALUES (?, ?, ?, ?, ?, ?)""",
|
||||
(file_hash, path, str(timestamp), engine, confidence, chunks),
|
||||
)
|
||||
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))
|
||||
|
||||
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))
|
||||
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))
|
||||
|
||||
def get_stats(self, days: int = 30) -> dict[str, Any]:
|
||||
cutoff = time.time() - days * 86400
|
||||
with self._conn() as conn:
|
||||
cursor = conn.execute("""
|
||||
SELECT
|
||||
COUNT(*) as total_processed,
|
||||
SUM(chunks_created) as total_chunks,
|
||||
AVG(confidence) as avg_confidence
|
||||
FROM processed_files
|
||||
""")
|
||||
row = cursor.fetchone()
|
||||
return dict(row) if row else {}
|
||||
|
||||
def start_run(self, date: str | None = None) -> int:
|
||||
if date is None:
|
||||
date = time.strftime("%Y-%m-%d")
|
||||
with self._conn() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT COALESCE(MAX(run_number), 0) + 1 FROM runs WHERE date = ?", (date,)
|
||||
)
|
||||
run_number = cursor.fetchone()[0]
|
||||
conn.execute(
|
||||
"INSERT INTO runs (date, run_number, total_images) VALUES (?, ?, 0)",
|
||||
(date, run_number),
|
||||
)
|
||||
return run_number
|
||||
|
||||
def complete_run(self, date: str, run_number: int, stats: dict[str, int]):
|
||||
with self._conn() as conn:
|
||||
conn.execute(
|
||||
"""UPDATE runs SET
|
||||
total_images = ?,
|
||||
successful = ?,
|
||||
failed = ?,
|
||||
chunks_created = ?,
|
||||
completed_at = ?
|
||||
WHERE date = ? AND run_number = ?""",
|
||||
(
|
||||
stats.get("total", 0),
|
||||
stats.get("successful", 0),
|
||||
stats.get("failed", 0),
|
||||
stats.get("chunks", 0),
|
||||
time.time(),
|
||||
date,
|
||||
run_number,
|
||||
),
|
||||
)
|
||||
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)}
|
||||
|
||||
|
||||
_db_instance: Database | None = None
|
||||
@@ -161,6 +117,5 @@ 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
|
||||
|
||||
@@ -7,9 +7,7 @@ from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import nbformat
|
||||
import platformdirs
|
||||
|
||||
from ocr_pipeline.config import settings
|
||||
from ocr_pipeline.utils.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
@@ -104,8 +102,11 @@ def migrate_notebook(notebook_path: Path, pipeline: Any) -> dict[str, Any]:
|
||||
continue
|
||||
|
||||
logger.info("migrating_image", path=str(img_path))
|
||||
pipeline.process_single(img_path)
|
||||
output_files = pipeline.process_single(img_path)
|
||||
if output_files:
|
||||
results["processed"] += 1
|
||||
else:
|
||||
results["skipped"] += 1
|
||||
|
||||
except Exception as e:
|
||||
logger.error("migration_failed", path=str(img_path), error=str(e))
|
||||
|
||||
@@ -1,144 +1,5 @@
|
||||
from __future__ import annotations
|
||||
"""Compatibility module for the canonical watcher package."""
|
||||
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from .watch.watcher import ScreenshotHandler, Watcher
|
||||
|
||||
from watchdog.events import FileSystemEventHandler
|
||||
from watchdog.observers import Observer
|
||||
|
||||
from ocr_pipeline.config import settings
|
||||
from ocr_pipeline.utils.db import ProcessedFile, get_db
|
||||
from ocr_pipeline.utils.logging import get_logger
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class ScreenshotHandler(FileSystemEventHandler):
|
||||
def __init__(self, process_callback: Callable[[Path], None]):
|
||||
self.process_callback = process_callback
|
||||
self.pending: dict[str, float] = {}
|
||||
self.lock = threading.Lock()
|
||||
self.debounce = settings.watch.debounce_seconds
|
||||
|
||||
def _should_process(self, path: Path) -> bool:
|
||||
if not path.is_file():
|
||||
return False
|
||||
if path.suffix.lower() not in (".png", ".jpg", ".jpeg", ".tiff", ".bmp"):
|
||||
return False
|
||||
for pattern in settings.watch.ignore_patterns:
|
||||
if path.match(pattern):
|
||||
return False
|
||||
return True
|
||||
|
||||
def on_created(self, event):
|
||||
if not event.is_directory:
|
||||
path = Path(event.src_path)
|
||||
if self._should_process(path):
|
||||
with self.lock:
|
||||
self.pending[str(path)] = time.time()
|
||||
|
||||
def on_modified(self, event):
|
||||
if not event.is_directory:
|
||||
path = Path(event.src_path)
|
||||
if self._should_process(path):
|
||||
with self.lock:
|
||||
self.pending[str(path)] = time.time()
|
||||
|
||||
def check_pending(self) -> list[Path]:
|
||||
now = time.time()
|
||||
to_process = []
|
||||
|
||||
with self.lock:
|
||||
for path_str, timestamp in list(self.pending.items()):
|
||||
if now - timestamp >= self.debounce:
|
||||
path = Path(path_str)
|
||||
if path.exists():
|
||||
to_process.append(path)
|
||||
del self.pending[path_str]
|
||||
|
||||
return to_process
|
||||
|
||||
|
||||
class Watcher:
|
||||
def __init__(self, process_callback: Callable[[Path], None]):
|
||||
self.process_callback = process_callback
|
||||
self.observer: Observer | None = None
|
||||
self.handler: ScreenshotHandler | None = None
|
||||
self.running = False
|
||||
self.db = get_db()
|
||||
self.run_id: int | None = None
|
||||
|
||||
def start(self) -> None:
|
||||
if self.running:
|
||||
return
|
||||
|
||||
self.run_id = self.db.start_run()
|
||||
self.running = True
|
||||
|
||||
self.handler = ScreenshotHandler(self._process_file)
|
||||
self.observer = Observer()
|
||||
|
||||
for base_path_str in settings.input.paths:
|
||||
base_path = Path(base_path_str).expanduser()
|
||||
if base_path.exists():
|
||||
self.observer.schedule(
|
||||
self.handler, str(base_path), recursive=settings.input.recursive
|
||||
)
|
||||
logger.info("watching_directory", path=str(base_path))
|
||||
|
||||
self.observer.start()
|
||||
|
||||
self._watch_thread = threading.Thread(target=self._watch_loop, daemon=True)
|
||||
self._watch_thread.start()
|
||||
|
||||
logger.info("watcher_started", run_id=self.run_id)
|
||||
|
||||
def _watch_loop(self) -> None:
|
||||
poll_interval = settings.watch.poll_interval
|
||||
processed = 0
|
||||
failed = 0
|
||||
|
||||
while self.running:
|
||||
time.sleep(poll_interval)
|
||||
|
||||
if self.handler:
|
||||
for path in self.handler.check_pending():
|
||||
try:
|
||||
self._process_file(path)
|
||||
processed += 1
|
||||
except Exception as e:
|
||||
logger.error("watch_process_failed", path=str(path), error=str(e))
|
||||
failed += 1
|
||||
|
||||
self.db.end_run(self.run_id, processed, failed)
|
||||
logger.info("watcher_stopped", run_id=self.run_id, processed=processed, failed=failed)
|
||||
|
||||
def _process_file(self, path: Path) -> None:
|
||||
from ocr_pipeline.ocr.engine import compute_image_hash
|
||||
|
||||
file_hash = compute_image_hash(path)
|
||||
if self.db.is_processed(path, file_hash):
|
||||
logger.debug("file_already_processed", path=str(path))
|
||||
return
|
||||
|
||||
logger.info("processing_new_file", path=str(path))
|
||||
self.process_callback(path)
|
||||
self.db.mark_processed(
|
||||
ProcessedFile(
|
||||
path=str(path),
|
||||
hash=file_hash,
|
||||
timestamp=time.time(),
|
||||
run_id=self.run_id,
|
||||
size=path.stat().st_size,
|
||||
mtime=path.stat().st_mtime,
|
||||
)
|
||||
)
|
||||
|
||||
def stop(self) -> None:
|
||||
self.running = False
|
||||
if self.observer:
|
||||
self.observer.stop()
|
||||
self.observer.join(timeout=5)
|
||||
logger.info("watcher_stopped")
|
||||
__all__ = ["ScreenshotHandler", "Watcher"]
|
||||
|
||||
@@ -74,7 +74,7 @@ class Watcher:
|
||||
if self.running:
|
||||
return
|
||||
|
||||
self.run_id = self.db.start_run()
|
||||
self.run_id, _ = self.db.start_run()
|
||||
self.running = True
|
||||
|
||||
self.handler = ScreenshotHandler(self._process_file)
|
||||
@@ -112,29 +112,14 @@ class Watcher:
|
||||
logger.error("watch_process_failed", path=str(path), error=str(e))
|
||||
failed += 1
|
||||
|
||||
self.db.complete_run(
|
||||
time.strftime("%Y-%m-%d"),
|
||||
self.run_id,
|
||||
{"total": processed, "successful": processed, "failed": failed, "chunks": 0},
|
||||
)
|
||||
if self.run_id is not None:
|
||||
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:
|
||||
from ocr_pipeline.ocr.engine import compute_image_hash
|
||||
|
||||
file_hash = compute_image_hash(path)
|
||||
if self.db.is_processed(file_hash, path):
|
||||
logger.debug("file_already_processed", path=str(path))
|
||||
return
|
||||
|
||||
logger.info("processing_new_file", path=str(path))
|
||||
# The pipeline owns deduplication and database writes, preventing double marks.
|
||||
self.process_callback(path)
|
||||
self.db.mark_processed(
|
||||
file_hash=file_hash,
|
||||
path=str(path),
|
||||
timestamp=time.time(),
|
||||
run_id=self.run_id,
|
||||
)
|
||||
|
||||
def stop(self) -> None:
|
||||
self.running = False
|
||||
|
||||
29
tests/test_output.py
Normal file
29
tests/test_output.py
Normal file
@@ -0,0 +1,29 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from ocr_pipeline.output.markdown import MarkdownWriter
|
||||
|
||||
|
||||
def test_writer_uses_metadata_and_unique_source_hash(tmp_path: Path, monkeypatch) -> None:
|
||||
from ocr_pipeline.output import markdown
|
||||
|
||||
monkeypatch.setattr(markdown.settings.output, "base_directory", str(tmp_path))
|
||||
monkeypatch.setattr(markdown.settings.output, "organize_by", "date_run")
|
||||
writer = MarkdownWriter()
|
||||
writer.start_run("2026-07-19T12:00:00+00:00", 2)
|
||||
metadata = {
|
||||
"source_path": "/screenshots/SCR-001.png",
|
||||
"source_hash": "0123456789abcdef",
|
||||
"timestamp": "2026-07-19T12:00:00+00:00",
|
||||
"chunk_index": 0,
|
||||
"total_chunks": 1,
|
||||
}
|
||||
|
||||
output = writer.write_chunk(content="recognized text", metadata=metadata)
|
||||
|
||||
assert output.name == "SCR-001_0123456789ab.md"
|
||||
text = output.read_text(encoding="utf-8")
|
||||
assert "# Screenshot: SCR-001.png" in text
|
||||
assert "recognized text" in text
|
||||
assert "source_hash: 0123456789abcdef" in text
|
||||
Reference in New Issue
Block a user