improve the OCR pipeline processing and outputs formatting

This commit is contained in:
2026-07-19 19:54:53 +02:00
parent d375db47c3
commit 1cf1c6eeab
30 changed files with 3838 additions and 3900 deletions

View File

@@ -14,8 +14,8 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
WORKDIR /app WORKDIR /app
COPY pyproject.toml ./ COPY pyproject.toml uv.lock ./
RUN pip install --no-cache-dir uv && uv sync --frozen RUN pip install --no-cache-dir uv && uv sync --frozen --no-dev
COPY src ./src COPY src ./src
COPY config.yaml ./ COPY config.yaml ./

View File

@@ -154,7 +154,7 @@ Lane 4: BRCA1 KO + pBRCA1-C61G mutant
Anti-BRCA1 (1:1000), Anti-β-actin (1:5000) Anti-BRCA1 (1:1000), Anti-β-actin (1:5000)
``` ```
Files land in `data/ocr_output/<date>/run_NNN/`. Files land in `data/ocr_output/<date>/run_NNN/`. Individual chunk files are retained for RAG indexing. Set `output.write_consolidated: true` to also write `all_ocr.md` containing one frontmatter block and all source text grouped by image.
## Life science specifics ## Life science specifics
@@ -169,6 +169,17 @@ First run downloads what it needs, cached in `~/.cache/ocr_pipeline/`:
- Table Transformer (~500MB) - Table Transformer (~500MB)
- LayoutParser PubLayNet (~300MB) - 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 ## TODO
- [ ] Build a knowledge graph from extracted entities/citations - [ ] Build a knowledge graph from extracted entities/citations

View File

@@ -8,6 +8,8 @@ input:
- "*.tiff" - "*.tiff"
- "*.bmp" - "*.bmp"
recursive: true recursive: true
exclude_patterns:
- "*.photoslibrary/*"
ocr: ocr:
engine: "paddleocr" engine: "paddleocr"
@@ -68,6 +70,8 @@ chunking:
keep_separator: true keep_separator: true
output: 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) # Uses platformdirs default: ~/.local/share/ocr-pipeline (Linux), ~/Library/Application Support/ocr-pipeline (macOS), %LOCALAPPDATA%\ocr-pipeline (Windows)
format: "markdown" format: "markdown"
organize_by: "date_run" organize_by: "date_run"

View File

@@ -1,3 +1,7 @@
[build-system]
requires = ["setuptools>=68"]
build-backend = "setuptools.build_meta"
[project] [project]
name = "ocr-pipeline" name = "ocr-pipeline"
version = "0.1.0" version = "0.1.0"
@@ -21,18 +25,22 @@ dependencies = [
"scikit-image>=0.22.0", "scikit-image>=0.22.0",
# OCR engines # OCR engines
"paddleocr>=2.7.0", "paddleocr>=2.7.0,<3.0.0",
"paddlepaddle>=2.6.0",
"pytesseract>=0.3.10", "pytesseract>=0.3.10",
# Table detection # Table detection
"torch>=2.2.0", "torch>=2.2.0",
"torchvision>=0.17.0", "torchvision>=0.17.0",
"transformers>=4.38.0", "transformers>=4.38.0",
"timm>=1.0.0",
"layoutparser>=0.3.0", "layoutparser>=0.3.0",
# NLP # NLP
"langchain-text-splitters>=0.0.2", "langchain-text-splitters>=0.0.2",
"spacy>=3.7.0", "spacy>=3.7.0",
"nbformat>=5.9.0",
"legacy-cgi>=2.6.2",
# Database # Database
"sqlite-utils>=3.37.0", "sqlite-utils>=3.37.0",
@@ -56,17 +64,17 @@ dev = [
"mypy>=1.8.0", "mypy>=1.8.0",
"pre-commit>=3.6.0", "pre-commit>=3.6.0",
] ]
grobid = [ grobid = []
"grobid-client>=0.8.0",
]
full = [ full = [
"ocr-pipeline[dev]", "ocr-pipeline[dev]",
"ocr-pipeline[grobid]",
] ]
[project.entry-points.console_scripts] [project.entry-points.console_scripts]
ocr-pipeline = "ocr_pipeline.cli:app" ocr-pipeline = "ocr_pipeline.cli:app"
[tool.setuptools.packages.find]
where = ["src"]
[tool.uv] [tool.uv]
dev-dependencies = [ dev-dependencies = [
"pytest>=8.0.0", "pytest>=8.0.0",

View File

@@ -1,15 +1,9 @@
from setuptools import setup, find_packages """Legacy setuptools compatibility shim.
setup( Project metadata, dependencies, and the src package layout are maintained solely in
name="kg_ocr", pyproject.toml. Install with `pip install .` or `uv sync`.
packages=find_packages(), """
python_requires=">=3.10",
install_requires=[ from setuptools import setup
"pytesseract",
"Pillow", setup()
"txtai",
"sentence-transformers",
"litellm",
"python-dotenv"
],
)

View File

@@ -15,19 +15,24 @@ Requires-Dist: opencv-python-headless>=4.9.0
Requires-Dist: pillow>=10.2.0 Requires-Dist: pillow>=10.2.0
Requires-Dist: numpy>=1.26.0 Requires-Dist: numpy>=1.26.0
Requires-Dist: scikit-image>=0.22.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: pytesseract>=0.3.10
Requires-Dist: torch>=2.2.0 Requires-Dist: torch>=2.2.0
Requires-Dist: torchvision>=0.17.0 Requires-Dist: torchvision>=0.17.0
Requires-Dist: transformers>=4.38.0 Requires-Dist: transformers>=4.38.0
Requires-Dist: timm>=1.0.0
Requires-Dist: layoutparser>=0.3.0 Requires-Dist: layoutparser>=0.3.0
Requires-Dist: langchain-text-splitters>=0.0.2 Requires-Dist: langchain-text-splitters>=0.0.2
Requires-Dist: spacy>=3.7.0 Requires-Dist: 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: sqlite-utils>=3.37.0
Requires-Dist: watchdog>=3.0.0 Requires-Dist: watchdog>=3.0.0
Requires-Dist: python-slugify>=8.0.0 Requires-Dist: python-slugify>=8.0.0
Requires-Dist: xxhash>=3.4.0 Requires-Dist: xxhash>=3.4.0
Requires-Dist: python-magic>=0.4.27 Requires-Dist: python-magic>=0.4.27
Requires-Dist: platformdirs>=4.2.0
Provides-Extra: dev Provides-Extra: dev
Requires-Dist: pytest>=8.0.0; extra == "dev" Requires-Dist: pytest>=8.0.0; extra == "dev"
Requires-Dist: pytest-cov>=4.1.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: mypy>=1.8.0; extra == "dev"
Requires-Dist: pre-commit>=3.6.0; extra == "dev" Requires-Dist: pre-commit>=3.6.0; extra == "dev"
Provides-Extra: grobid Provides-Extra: grobid
Requires-Dist: grobid-client>=0.8.0; extra == "grobid"
Provides-Extra: full Provides-Extra: full
Requires-Dist: ocr-pipeline[dev]; extra == "full" Requires-Dist: ocr-pipeline[dev]; extra == "full"
Requires-Dist: ocr-pipeline[grobid]; extra == "full"
# OCR Pipeline for Life Science Screenshots # 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 ## Quick start
- **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
```bash ```bash
# Install with uv (recommended) uv sync # or: pip install -e .
uv sync
# Or with pip
pip install -e .
# Install pre-commit hooks
pre-commit install pre-commit install
# Run on your screenshots ocr-pipeline run # process screenshots
ocr-pipeline run ocr-pipeline watch # watch a folder
ocr-pipeline reprocess ocr_sc.ipynb # pull screenshots out of a notebook
# Watch for new screenshots
ocr-pipeline watch
# Reprocess from notebook
ocr-pipeline reprocess ocr_sc.ipynb
# Check status
ocr-pipeline status 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 ```yaml
input: input:
@@ -120,11 +113,11 @@ detectors:
entities: entities:
enabled: true enabled: true
model: "en_core_sci_lg" # scispaCy large model model: "en_core_sci_lg"
citations: citations:
regex_enabled: true 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: chunking:
chunk_size: 1000 chunk_size: 1000
@@ -153,9 +146,11 @@ watch:
db_path: "./data/processed_files.db" 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 ```markdown
--- ---
@@ -178,9 +173,9 @@ total_chunks: 2
## Figures ## Figures
### Figure 1 ### Figure 1
- **BBox**: [100, 200, 800, 600] - BBox: [100, 200, 800, 600]
- **Confidence**: 0.92 - Confidence: 0.92
- **Caption**: "Western blot showing BRCA1 expression..." - Caption: "Western blot showing BRCA1 expression..."
## Detected Entities ## Detected Entities
@@ -205,72 +200,33 @@ Lane 4: BRCA1 KO + pBRCA1-C61G mutant
Anti-BRCA1 (1:1000), Anti-β-actin (1:5000) Anti-BRCA1 (1:1000), Anti-β-actin (1:5000)
``` ```
## Directory Structure Files land in `data/ocr_output/<date>/run_NNN/`.
``` ## Life science specifics
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/
└── ...
```
## 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 ## Models
# Build
docker build -t ocr-pipeline .
# Run once First run downloads what it needs, cached in `~/.cache/ocr_pipeline/`:
docker run -v ~/Pictures:/data/screenshots -v ./data:/app/data ocr-pipeline run
# Watch mode - PaddleOCR models (~200MB)
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)
- scispaCy `en_core_sci_lg` (~800MB) - scispaCy `en_core_sci_lg` (~800MB)
- Table Transformer (~500MB) - Table Transformer (~500MB)
- LayoutParser PubLayNet (~300MB) - 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

View File

@@ -1,5 +1,6 @@
README.md README.md
pyproject.toml pyproject.toml
setup.py
src/ocr_pipeline/__init__.py src/ocr_pipeline/__init__.py
src/ocr_pipeline/cli.py src/ocr_pipeline/cli.py
src/ocr_pipeline/config.py src/ocr_pipeline/config.py
@@ -20,6 +21,7 @@ src/ocr_pipeline/detectors/tables.py
src/ocr_pipeline/ocr/__init__.py src/ocr_pipeline/ocr/__init__.py
src/ocr_pipeline/ocr/engine.py src/ocr_pipeline/ocr/engine.py
src/ocr_pipeline/ocr/parallel.py src/ocr_pipeline/ocr/parallel.py
src/ocr_pipeline/ocr/preprocess.py
src/ocr_pipeline/output/__init__.py src/ocr_pipeline/output/__init__.py
src/ocr_pipeline/output/markdown.py src/ocr_pipeline/output/markdown.py
src/ocr_pipeline/postprocess/__init__.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/utils/migrate.py
src/ocr_pipeline/watch/__init__.py src/ocr_pipeline/watch/__init__.py
src/ocr_pipeline/watch/watcher.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 tests/test_postprocess.py

View File

@@ -9,19 +9,24 @@ opencv-python-headless>=4.9.0
pillow>=10.2.0 pillow>=10.2.0
numpy>=1.26.0 numpy>=1.26.0
scikit-image>=0.22.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 pytesseract>=0.3.10
torch>=2.2.0 torch>=2.2.0
torchvision>=0.17.0 torchvision>=0.17.0
transformers>=4.38.0 transformers>=4.38.0
timm>=1.0.0
layoutparser>=0.3.0 layoutparser>=0.3.0
langchain-text-splitters>=0.0.2 langchain-text-splitters>=0.0.2
spacy>=3.7.0 spacy>=3.7.0
nbformat>=5.9.0
legacy-cgi>=2.6.2
sqlite-utils>=3.37.0 sqlite-utils>=3.37.0
watchdog>=3.0.0 watchdog>=3.0.0
python-slugify>=8.0.0 python-slugify>=8.0.0
xxhash>=3.4.0 xxhash>=3.4.0
python-magic>=0.4.27 python-magic>=0.4.27
platformdirs>=4.2.0
[dev] [dev]
pytest>=8.0.0 pytest>=8.0.0
@@ -33,7 +38,5 @@ pre-commit>=3.6.0
[full] [full]
ocr-pipeline[dev] ocr-pipeline[dev]
ocr-pipeline[grobid]
[grobid] [grobid]
grobid-client>=0.8.0

View File

@@ -1,191 +1,16 @@
from __future__ import annotations """Backward-compatible citation exports.
import re The canonical implementation lives in ``regex_extractor``. Keeping this module as a
from dataclasses import dataclass thin re-export avoids two incompatible Citation models and removes an undeclared
from typing import Any ``requests`` dependency from the runtime path.
"""
import requests from .regex_extractor import (
Citation,
from ocr_pipeline.config import settings GrobidClient,
from ocr_pipeline.utils.logging import get_logger RegexCitationExtractor,
extract_citations,
logger = get_logger(__name__) format_citations_for_markdown,
@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,
) )
@classmethod __all__ = ["Citation", "GrobidClient", "RegexCitationExtractor", "extract_citations", "format_citations_for_markdown"]
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)

View File

@@ -2,8 +2,8 @@ from __future__ import annotations
import re import re
from dataclasses import dataclass from dataclasses import dataclass
from urllib.parse import urlencode
import httpx from urllib.request import Request, urlopen
from ocr_pipeline.config import settings from ocr_pipeline.config import settings
from ocr_pipeline.utils.logging import get_logger from ocr_pipeline.utils.logging import get_logger
@@ -26,7 +26,7 @@ class RegexCitationExtractor:
re.IGNORECASE, re.IGNORECASE,
) )
PMID_PATTERN = re.compile( 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, re.IGNORECASE,
) )
ARXIV_PATTERN = re.compile( ARXIV_PATTERN = re.compile(
@@ -49,7 +49,7 @@ class RegexCitationExtractor:
Citation( Citation(
raw=match.group(), raw=match.group(),
type="doi", type="doi",
identifier=match.group().lower(), identifier=match.group().rstrip(".,;:)]}").lower(),
confidence=0.95, confidence=0.95,
context=text[ctx_start:ctx_end].strip(), context=text[ctx_start:ctx_end].strip(),
) )
@@ -108,13 +108,14 @@ class GrobidClient:
return [] return []
try: try:
with httpx.Client(timeout=self.timeout) as client: payload = urlencode({"text": text}).encode("utf-8")
response = client.post( request = Request(
f"{self.url}/api/processCitationText", f"{self.url}/api/processCitationText",
data={"text": text}, data=payload,
headers={"Content-Type": "application/x-www-form-urlencoded"},
) )
response.raise_for_status() with urlopen(request, timeout=self.timeout) as response:
return self._parse_grobid_xml(response.text) return self._parse_grobid_xml(response.read().decode("utf-8"))
except Exception as e: except Exception as e:
logger.warning("grobid_request_failed", error=str(e)) logger.warning("grobid_request_failed", error=str(e))
return [] return []

View File

@@ -44,6 +44,8 @@ def run(
engine: str | None = typer.Option( engine: str | None = typer.Option(
None, "--engine", "-e", help="OCR engine (paddleocr/tesseract/auto)" None, "--engine", "-e", help="OCR engine (paddleocr/tesseract/auto)"
), ),
exclude: list[str] = typer.Option([], "--exclude", help="Glob pattern to exclude; repeatable"),
force: bool = typer.Option(False, "--force", help="Reprocess files already recorded as processed"),
): ):
"""Run OCR pipeline on all screenshots""" """Run OCR pipeline on all screenshots"""
if input_dir: if input_dir:
@@ -54,6 +56,10 @@ def run(
settings.processing.workers = workers settings.processing.workers = workers
if engine: if engine:
settings.ocr.engine = engine settings.ocr.engine = engine
if exclude:
settings.input.exclude_patterns = exclude
if force:
settings.processing.skip_existing = False
pipeline = OCRPipeline() pipeline = OCRPipeline()
result = pipeline.process_all() 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")): def stats(days: int = typer.Option(30, "--days", "-d", help="Number of days to analyze")):
"""Show processing statistics""" """Show processing statistics"""
pipeline = OCRPipeline() pipeline = OCRPipeline()
stats = pipeline.get_stats() stats = pipeline.get_stats(days)
table = Table(title=f"Statistics (last {days} days)") table = Table(title=f"Statistics (last {days} days)")
table.add_column("Metric", style="cyan") table.add_column("Metric", style="cyan")

View File

@@ -3,14 +3,13 @@ from __future__ import annotations
import os import os
import platform import platform
from pathlib import Path from pathlib import Path
from typing import Any, ClassVar from typing import Any, ClassVar, Literal
import platformdirs import platformdirs
import yaml import yaml
from pydantic import Field, field_validator from pydantic import Field, field_validator
from pydantic_settings import BaseSettings, SettingsConfigDict from pydantic_settings import BaseSettings, SettingsConfigDict
APP_NAME = "ocr-pipeline" APP_NAME = "ocr-pipeline"
APP_AUTHOR = "aman" APP_AUTHOR = "aman"
@@ -43,6 +42,7 @@ class InputConfig(BaseSettings):
default_factory=lambda: ["SCR-*.png", "*.jpg", "*.jpeg", "*.tiff", "*.bmp"] default_factory=lambda: ["SCR-*.png", "*.jpg", "*.jpeg", "*.tiff", "*.bmp"]
) )
recursive: bool = True recursive: bool = True
exclude_patterns: list[str] = Field(default_factory=lambda: ["*.photoslibrary/*"])
@field_validator("paths", mode="before") @field_validator("paths", mode="before")
@classmethod @classmethod
@@ -60,7 +60,7 @@ class PreprocessConfig(BaseSettings):
class OCRConfig(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"]) languages: list[str] = Field(default_factory=lambda: ["en", "latin"])
use_gpu: bool = False use_gpu: bool = False
use_angle_cls: bool = True use_angle_cls: bool = True
@@ -73,32 +73,32 @@ class OCRConfig(BaseSettings):
class ProcessingConfig(BaseSettings): class ProcessingConfig(BaseSettings):
workers: int = 4 workers: int = Field(default=4, ge=1)
batch_size: int = 10 batch_size: int = Field(default=10, ge=1)
retry_attempts: int = 2 retry_attempts: int = Field(default=2, ge=0)
timeout_per_image: int = 120 timeout_per_image: int = Field(default=120, ge=1)
skip_existing: bool = True skip_existing: bool = True
class DetectorConfig(BaseSettings): class DetectorConfig(BaseSettings):
enabled: bool = True enabled: bool = True
model: str = "" model: str = ""
confidence_threshold: float = 0.7 confidence_threshold: float = Field(default=0.7, ge=0.0, le=1.0)
class FigureDetectorConfig(DetectorConfig): class FigureDetectorConfig(DetectorConfig):
model: str = "lp://PubLayNet/faster_rcnn_R_50_FPN_3x/config" 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): class TableDetectorConfig(DetectorConfig):
model: str = "microsoft/table-transformer-detection" 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): class CaptionDetectorConfig(DetectorConfig):
model: str = "lp://PubLayNet/faster_rcnn_R_50_FPN_3x/config" 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): class DetectorsConfig(BaseSettings):
@@ -149,6 +149,7 @@ class FrontmatterConfig(BaseSettings):
"ocr_confidence_mean", "ocr_confidence_mean",
"language", "language",
"detected_entities", "detected_entities",
"entity_extraction_backend",
"has_figures", "has_figures",
"has_tables", "has_tables",
"citations_found", "citations_found",
@@ -160,22 +161,24 @@ class FrontmatterConfig(BaseSettings):
class OutputConfig(BaseSettings): class OutputConfig(BaseSettings):
write_consolidated: bool = False
consolidated_filename: str = "all_ocr.md"
base_directory: str = Field( base_directory: str = Field(
default_factory=lambda: str(platformdirs.user_data_dir(APP_NAME, APP_AUTHOR)) default_factory=lambda: str(platformdirs.user_data_dir(APP_NAME, APP_AUTHOR))
) )
format: str = "markdown" 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) frontmatter: FrontmatterConfig = Field(default_factory=FrontmatterConfig)
class WatchConfig(BaseSettings): class WatchConfig(BaseSettings):
enabled: bool = True 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) ignore_patterns: list[str] = Field(default_factory=_get_platform_ignore_patterns)
db_path: str = Field( db_path: str = Field(
default_factory=lambda: str(Path(platformdirs.user_data_dir(APP_NAME, APP_AUTHOR)) / "processed_files.db") default_factory=lambda: str(Path(platformdirs.user_data_dir(APP_NAME, APP_AUTHOR)) / "processed_files.db")
) )
poll_interval: float = 1.0 poll_interval: float = Field(default=1.0, gt=0)
class _BaseSettings(BaseSettings): class _BaseSettings(BaseSettings):

View File

@@ -33,7 +33,10 @@ class FigureDetector:
try: try:
import layoutparser as lp 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, config_path=settings.detectors.figures.model,
label_map={0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"}, label_map={0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"},
extra_config=[ extra_config=[
@@ -67,7 +70,10 @@ class FigureDetector:
block.block.x_2, block.block.x_2,
block.block.y_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( figures.append(
DetectedFigure( DetectedFigure(
bbox=[float(x1), float(y1), float(x2), float(y2)], bbox=[float(x1), float(y1), float(x2), float(y2)],
@@ -98,7 +104,10 @@ class CaptionExtractor:
try: try:
import layoutparser as lp 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, config_path=settings.detectors.captions.model,
label_map={0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"}, label_map={0: "Text", 1: "Title", 2: "List", 3: "Table", 4: "Figure"},
extra_config=[ extra_config=[
@@ -156,11 +165,14 @@ class CaptionExtractor:
return figures return figures
def detect_figures(image: np.ndarray) -> list[DetectedFigure]: _figure_detector = FigureDetector()
detector = FigureDetector() _caption_extractor = CaptionExtractor()
figures = detector.detect(image)
captioner = CaptionExtractor()
def detect_figures(image: np.ndarray) -> list[DetectedFigure]:
figures = _figure_detector.detect(image)
captioner = _caption_extractor
figures = captioner.extract_captions(image, figures) figures = captioner.extract_captions(image, figures)
return figures return figures

View File

@@ -59,6 +59,7 @@ class TableDetector:
try: try:
import torch import torch
original_height, original_width = image.shape[:2]
pil_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) pil_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
pil_image = cv2.resize(pil_image, (800, 800)) pil_image = cv2.resize(pil_image, (800, 800))
@@ -77,8 +78,12 @@ class TableDetector:
for score, _label, box in zip( for score, _label, box in zip(
results["scores"], results["labels"], results["boxes"], strict=True 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, 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] crop = image[y1:y2, x1:x2]
tables.append( tables.append(
DetectedTable( DetectedTable(
@@ -122,6 +127,7 @@ class TableStructureRecognizer:
self._initialized = True self._initialized = True
def recognize(self, table_crop: np.ndarray) -> tuple[list[dict[str, Any]] | None, str | None]: 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: if self._model is None or self._processor is None:
return None, None return None, None
@@ -150,8 +156,9 @@ class TableStructureRecognizer:
} }
) )
markdown = self._structure_to_markdown(structure, table_crop.shape) # Structure models locate rows/cells but do not provide cell text. Do not
return structure, markdown # emit labels as fabricated table content; OCR must be associated separately.
return structure, None
except Exception as e: except Exception as e:
logger.warning("table_structure_recognition_failed", error=str(e)) logger.warning("table_structure_recognition_failed", error=str(e))
return None, None return None, None
@@ -185,17 +192,19 @@ class TableStructureRecognizer:
return None return None
_table_detector = TableDetector()
_table_structure_recognizer = TableStructureRecognizer()
def detect_tables(image: np.ndarray) -> list[DetectedTable]: def detect_tables(image: np.ndarray) -> list[DetectedTable]:
detector = TableDetector() tables = _table_detector.detect(image)
tables = detector.detect(image)
if not tables: if not tables:
return tables return tables
recognizer = TableStructureRecognizer()
for table in tables: for table in tables:
if table.image_crop is not None: 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.structure = structure
table.markdown = markdown table.markdown = markdown

View File

@@ -63,6 +63,7 @@ class OCRResult:
boxes: list[dict[str, Any]] boxes: list[dict[str, Any]]
raw_result: Any raw_result: Any
processing_time: float processing_time: float
fallback_reason: str | None = None
class PaddleOCREngine: class PaddleOCREngine:
@@ -78,17 +79,23 @@ class PaddleOCREngine:
paddle_lang = map_languages(settings.ocr.languages, "paddleocr") paddle_lang = map_languages(settings.ocr.languages, "paddleocr")
self._ocr = PaddleOCR( legacy_options = {
use_angle_cls=settings.ocr.use_angle_cls, "use_angle_cls": settings.ocr.use_angle_cls,
lang=paddle_lang, "lang": paddle_lang,
use_gpu=settings.ocr.use_gpu, "use_gpu": settings.ocr.use_gpu,
show_log=False, "show_log": False,
det_db_thresh=settings.ocr.det_db_thresh, "det_db_thresh": settings.ocr.det_db_thresh,
det_db_box_thresh=settings.ocr.det_db_box_thresh, "det_db_box_thresh": settings.ocr.det_db_box_thresh,
det_db_unclip_ratio=settings.ocr.det_db_unclip_ratio, "det_db_unclip_ratio": settings.ocr.det_db_unclip_ratio,
rec_batch_num=settings.ocr.rec_batch_num, "rec_batch_num": settings.ocr.rec_batch_num,
cpu_threads=settings.ocr.cpu_threads, "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 self._initialized = True
logger.info("paddleocr_initialized", language=paddle_lang) logger.info("paddleocr_initialized", language=paddle_lang)
except Exception as e: except Exception as e:

View File

@@ -26,7 +26,8 @@ class ProcessedImage:
def _process_single_image(args: tuple[Path, dict]) -> ProcessedImage: def _process_single_image(args: tuple[Path, dict]) -> ProcessedImage:
"""Worker function for multiprocessing - must be at module level.""" """Worker function for multiprocessing - must be at module level."""
path, config_dict = args 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() preprocessor = ImagePreprocessor()
ocr_engine = OCREngine() ocr_engine = OCREngine()
@@ -86,7 +87,13 @@ class ParallelProcessor:
unique_images = [] unique_images = []
seen = set() seen = set()
for img in all_images: 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: if img not in seen:
seen.add(img) seen.add(img)
unique_images.append(img) unique_images.append(img)

View File

@@ -1,119 +1,95 @@
from __future__ import annotations from __future__ import annotations
import hashlib import hashlib
from pathlib import Path
import cv2 import cv2
import numpy as np import numpy as np
from ocr_pipeline.config import settings from ocr_pipeline.config import settings
from ocr_pipeline.utils.logging import get_logger
logger = get_logger(__name__)
class ImagePreprocessor: 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.config = settings.ocr.preprocess
self.max_dim = self.config.max_dimension self.max_dim = self.config.max_dimension
def load_image(self, path: str) -> np.ndarray: def load_image(self, path: str | Path) -> np.ndarray:
"""Load image from file path."""
image = cv2.imread(str(path)) image = cv2.imread(str(path))
if image is None: if image is None:
raise ValueError(f"Failed to load image: {path}") raise ValueError(f"Failed to load image: {path}")
return image return image
def resize_if_needed(self, image: np.ndarray) -> np.ndarray: def resize_if_needed(self, image: np.ndarray) -> np.ndarray:
"""Resize image if it exceeds max dimension.""" height, width = image.shape[:2]
h, w = image.shape[:2] if max(height, width) <= self.max_dim:
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)
return image 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: def deskew(self, image: np.ndarray) -> np.ndarray:
"""Correct skew in the image."""
if not self.config.deskew: if not self.config.deskew:
return image return image
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
coords = np.column_stack(np.where(gray > 0)) _, foreground = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
if len(coords) == 0: coords = np.column_stack(np.where(foreground > 0))
if len(coords) < 20:
return image return image
angle = cv2.minAreaRect(coords)[-1] angle = cv2.minAreaRect(coords)[-1]
if angle < -45: angle = -(90 + angle) if angle < -45 else -angle
angle = -(90 + angle) if abs(angle) <= 0.5:
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)
return image 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: def denoise(self, image: np.ndarray) -> np.ndarray:
"""Remove noise from image.""" return cv2.fastNlMeansDenoisingColored(image, None, 10, 10, 7, 21) if self.config.denoise else image
if not self.config.denoise:
return image
return cv2.fastNlMeansDenoisingColored(image, None, 10, 10, 7, 21)
def apply_clahe(self, image: np.ndarray) -> np.ndarray: def apply_clahe(self, image: np.ndarray) -> np.ndarray:
"""Apply CLAHE for contrast enhancement."""
if not self.config.clahe: if not self.config.clahe:
return image return image
lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB) lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
l, a, b = cv2.split(lab) lightness, a_channel, b_channel = cv2.split(lab)
clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)) lightness = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)).apply(lightness)
l = clahe.apply(l) return cv2.cvtColor(cv2.merge((lightness, a_channel, b_channel)), cv2.COLOR_LAB2BGR)
lab = cv2.merge((l, a, b))
return cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)
def remove_lines(self, image: np.ndarray) -> np.ndarray: def remove_lines(self, image: np.ndarray) -> np.ndarray:
"""Remove horizontal and vertical lines (tables, grids)."""
if not self.config.remove_lines: if not self.config.remove_lines:
return image return image
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
_, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU) _, binary = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
horizontal_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (40, 1)) horizontal = cv2.morphologyEx(binary, cv2.MORPH_OPEN, cv2.getStructuringElement(cv2.MORPH_RECT, (40, 1)))
vertical_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 40)) vertical = cv2.morphologyEx(binary, cv2.MORPH_OPEN, cv2.getStructuringElement(cv2.MORPH_RECT, (1, 40)))
horizontal_lines = cv2.morphologyEx(binary, cv2.MORPH_OPEN, horizontal_kernel) mask = cv2.bitwise_or(horizontal, vertical)
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)
result = image.copy() result = image.copy()
result[mask == 0] = [255, 255, 255] result[mask > 0] = (255, 255, 255)
return result return result
def adaptive_threshold(self, image: np.ndarray) -> np.ndarray: def adaptive_threshold(self, image: np.ndarray) -> np.ndarray:
"""Apply adaptive thresholding for binarization."""
if not self.config.adaptive_threshold: if not self.config.adaptive_threshold:
return image return image
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
binary = cv2.adaptiveThreshold( binary = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2)
gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
cv2.THRESH_BINARY, 11, 2
)
return cv2.cvtColor(binary, cv2.COLOR_GRAY2BGR) return cv2.cvtColor(binary, cv2.COLOR_GRAY2BGR)
def preprocess(self, path: str) -> np.ndarray: def preprocess_image(self, image: np.ndarray) -> np.ndarray:
"""Full preprocessing pipeline."""
image = self.load_image(path)
image = self.resize_if_needed(image) image = self.resize_if_needed(image)
image = self.deskew(image) image = self.deskew(image)
image = self.denoise(image) image = self.denoise(image)
image = self.apply_clahe(image) image = self.apply_clahe(image)
image = self.remove_lines(image) image = self.remove_lines(image)
image = self.adaptive_threshold(image) return self.adaptive_threshold(image)
return image
def preprocess(self, path: str | Path) -> np.ndarray:
return self.preprocess_image(self.load_image(path))
def compute_image_hash(path: str) -> str: def compute_image_hash(path: str | Path) -> str:
"""Compute SHA256 hash of image file for deduplication."""
hasher = hashlib.sha256() hasher = hashlib.sha256()
with open(path, "rb") as f: with Path(path).open("rb") as handle:
for chunk in iter(lambda: f.read(8192), b""): for chunk in iter(lambda: handle.read(8192), b""):
hasher.update(chunk) hasher.update(chunk)
return hasher.hexdigest() return hasher.hexdigest()

View File

@@ -3,13 +3,19 @@ from __future__ import annotations
import datetime import datetime
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import Any from typing import TYPE_CHECKING, Any
import yaml import yaml
from ocr_pipeline.config import settings from ocr_pipeline.config import settings
from ocr_pipeline.utils.logging import get_logger from ocr_pipeline.utils.logging import get_logger
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__) logger = get_logger(__name__)
@@ -30,131 +36,144 @@ class OutputMetadata:
class MarkdownWriter: class MarkdownWriter:
def __init__(self): def __init__(self) -> None:
self.base_dir = Path(settings.output.base_directory).expanduser().resolve() self.base_dir = Path(settings.output.base_directory).expanduser().resolve()
self.organize_by = settings.output.organize_by self.organize_by = settings.output.organize_by
self.frontmatter_fields = settings.output.frontmatter.fields 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")) dt = datetime.datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
if self.organize_by == "date_run": if self.organize_by == "date_run":
date_dir = self.base_dir / dt.strftime("%Y-%m-%d") self._run_dir = self.base_dir / dt.strftime("%Y-%m-%d") / f"run_{run_number:03d}"
run_dirs = sorted(date_dir.glob("run_*")) elif self.organize_by == "flat":
run_num = len(run_dirs) + 1 self._run_dir = self.base_dir
return date_dir / f"run_{run_num:03d}"
elif self.organize_by == "source_dir":
return self.base_dir / "by_source"
else: 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 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( def write_chunk(
self, self,
*,
content: str, content: str,
metadata: dict[str, Any], metadata: dict[str, Any],
figures: list[Any] = None, figures: list[DetectedFigure] | None = None,
tables: list[Any] = None, tables: list[DetectedTable] | None = None,
entities: Any = None, entities: EntityResult | None = None,
citations: list[Any] = None, citations: list[Citation] | None = None,
) -> Path: ) -> Path:
output_dir = self._get_output_dir(metadata.get("timestamp", "")) output_dir = self._get_output_dir(metadata)
output_dir.mkdir(parents=True, exist_ok=True) 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_path = Path(metadata.get("source_path", "unknown")) source_hash = str(metadata.get("source_hash", ""))[:12]
safe_name = "".join(c for c in source_path.stem if c.isalnum() or c in "-_")[:80] chunk_idx = int(metadata.get("chunk_index", 0))
chunk_idx = metadata.get("chunk_index", 0) total_chunks = int(metadata.get("total_chunks", 1))
total_chunks = 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"
if total_chunks > 1:
filename = f"{safe_name}_chunk_{chunk_idx:03d}.md"
else:
filename = f"{safe_name}.md"
output_path = output_dir / filename output_path = output_dir / filename
frontmatter = self._build_frontmatter(metadata) frontmatter = self._build_frontmatter(metadata)
body = self._build_body(content, figures, tables, entities, citations) body = self._build_body(content, source_path.name, figures or [], tables or [], entities, citations or [])
self._written_chunks.append((dict(metadata), content))
with open(output_path, "w") as f: with output_path.open("w", encoding="utf-8") as handle:
f.write("---\n") handle.write("---\n")
yaml.dump(frontmatter, f, sort_keys=False, allow_unicode=True) yaml.safe_dump(frontmatter, handle, sort_keys=False, allow_unicode=True)
f.write("---\n\n") handle.write("---\n\n")
f.write(body) handle.write(body)
logger.debug("markdown_written", path=str(output_path)) logger.debug("markdown_written", path=str(output_path))
return 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]: def _build_frontmatter(self, metadata: dict[str, Any]) -> dict[str, Any]:
fm = {} return {field: metadata[field] for field in self.frontmatter_fields if field in metadata}
for field in self.frontmatter_fields:
if field in metadata:
fm[field] = metadata[field]
return fm
def _build_body( def _build_body(
self, self,
content: str, content: str,
figures: list[Any] = None, source_name: str,
tables: list[Any] = None, figures: list[DetectedFigure],
entities: Any = None, tables: list[DetectedTable],
citations: list[Any] = None, entities: EntityResult | None,
citations: list[Citation],
) -> str: ) -> str:
parts = [] parts = [f"# Screenshot: {source_name}", "", "## Extracted Text", "", content.strip(), ""]
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("")
if figures: 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("") 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: if tables:
parts.append("## Tables Detected") parts.extend(["## Tables Detected", ""])
parts.append("") for index, table in enumerate(tables, 1):
for i, table in enumerate(tables): parts.extend([f"### Table {index}", f"- **Bounding Box**: {table.bbox}", f"- **Confidence**: {table.confidence:.2f}"])
parts.append(f"### Table {i + 1}")
parts.append(f"- **Bounding Box**: {table.bbox}")
parts.append(f"- **Confidence**: {table.confidence:.2f}")
if table.markdown: if table.markdown:
parts.extend(["", "#### Table Content (Markdown)", "", table.markdown])
parts.append("") parts.append("")
parts.append("#### Table Content (Markdown)")
parts.append("")
parts.append(table.markdown)
parts.append("")
if entities and entities.entities: 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("")
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: 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("")
parts.append("| Type | Identifier |")
parts.append("|------|------------|")
for cit in citations:
parts.append(f"| {cit.type.upper()} | `{cit.identifier}` |")
parts.append("")
return "\n".join(parts) return "\n".join(parts)

View File

@@ -1,7 +1,7 @@
from __future__ import annotations from __future__ import annotations
import datetime
import time import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from dataclasses import dataclass, field from dataclasses import dataclass, field
from pathlib import Path from pathlib import Path
from typing import Any 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.ocr.parallel import ParallelProcessor
from ocr_pipeline.output import MarkdownWriter from ocr_pipeline.output import MarkdownWriter
from ocr_pipeline.postprocess import chunk_text, clean_text, extract_entities 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.utils.logging import get_logger
from ocr_pipeline.watch import Watcher from ocr_pipeline.watch import Watcher
try:
from rich.console import Console
_console = Console()
except ImportError:
_console = None
logger = get_logger(__name__) logger = get_logger(__name__)
@@ -32,157 +25,104 @@ class PipelineResult:
total_images: int = 0 total_images: int = 0
successful: int = 0 successful: int = 0
failed: int = 0 failed: int = 0
skipped: int = 0
chunks_created: int = 0 chunks_created: int = 0
output_files: list[Path] = field(default_factory=list) 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: class OCRPipeline:
def __init__(self): def __init__(self) -> None:
self.engine = OCREngine() self.engine = OCREngine()
self.preprocessor = ImagePreprocessor() self.preprocessor = ImagePreprocessor()
self.writer = MarkdownWriter() self.writer = MarkdownWriter()
self.db = get_db() self.db: Database = get_db()
self._run_id: int | None = None
def _process_single(self, image_path: Path) -> list[dict[str, Any]]:
image = self.preprocessor.preprocess(image_path)
result = self.engine.process(image)
def _process_single(self, image_path: Path) -> list[Path]:
image_hash = compute_image_hash(image_path) image_hash = compute_image_hash(image_path)
if self.db.is_processed(image_hash, image_path): if settings.processing.skip_existing and self.db.is_processed(image_hash, image_path):
logger.debug("already_processed", path=str(image_path)) logger.info("already_processed", path=str(image_path))
return [] 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) chunks = chunk_text(cleaned.text)
entities = extract_entities(cleaned.text) entities = extract_entities(cleaned.text)
figures = detect_figures(image) # Detectors receive the original image: binarization/line removal damages their inputs.
tables = detect_tables(image) figures = detect_figures(original_image)
tables = detect_tables(original_image)
citations = extract_citations(cleaned.text) citations = extract_citations(cleaned.text)
timestamp = datetime.datetime.now(datetime.UTC).isoformat()
output_files = [] output_files: list[Path] = []
for i, chunk in enumerate(chunks): for chunk in chunks:
output_path = self.writer.write_chunk( metadata = {
source_path=image_path, "source_path": str(image_path), "source_hash": image_hash, "timestamp": timestamp,
source_hash=image_hash, "ocr_engine": ocr_result.engine, "ocr_confidence_mean": ocr_result.confidence,
text=chunk.content, "language": ocr_result.language, "entity_extraction_backend": entities.backend, "detected_entities": [entity.text for entity in entities.entities],
entities=[e.text for e in entities.entities], "has_figures": bool(figures), "has_tables": bool(tables),
figures=figures, "citations_found": [f"{citation.type}:{citation.identifier}" for citation in citations],
tables=tables, "chunk_index": chunk.chunk_index, "total_chunks": len(chunks),
citations=[f"{c.type}:{c.identifier}" for c in citations], }
chunk_index=i, output_files.append(self.writer.write_chunk(content=chunk.content, metadata=metadata, figures=figures, tables=tables, entities=entities, citations=citations))
total_chunks=len(chunks),
ocr_engine=result.engine,
ocr_confidence=result.confidence,
language=result.language,
)
output_files.append(output_path)
self.db.mark_processed( 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)
file_hash=image_hash, return output_files
path=str(image_path),
timestamp=time.time(),
engine=result.engine,
confidence=result.confidence,
chunks=len(chunks),
run_id=0,
)
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: def process_all(self, image_paths: list[Path] | None = None) -> PipelineResult:
"""Process images through the full pipeline."""
if image_paths is None: if image_paths is None:
processor = ParallelProcessor() image_paths = ParallelProcessor().find_images()
results = processor.process()
image_paths = [r.path for r in results if r.success]
else: else:
# Validate paths exist image_paths = [path for path in image_paths if path.exists()]
image_paths = [p for p in image_paths if p.exists()]
result = PipelineResult(total_images=len(image_paths)) result = PipelineResult(total_images=len(image_paths))
if not image_paths: if not image_paths:
return result return result
with ProcessPoolExecutor(max_workers=settings.processing.workers) as executor: run_timestamp = datetime.datetime.now(datetime.UTC).isoformat()
future_to_path = {executor.submit(self._process_single, p): p for p in image_paths} self._run_id, run_number = self.db.start_run()
self.writer.start_run(run_timestamp, run_number)
for future in as_completed(future_to_path):
path = future_to_path[future]
try: 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.successful += 1
result.chunks_created += sum(info["chunks"] for info in output_info) result.chunks_created += len(outputs)
result.output_files.extend([info["path"] for info in output_info]) result.output_files.extend(outputs)
except Exception as e: except Exception as exc:
result.failed += 1 result.failed += 1
result.errors.append({"path": str(path), "error": str(e)}) result.errors.append({"path": str(path), "error": str(exc)})
logger.error("pipeline_failed", path=str(path), error=str(e)) 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 return result
def watch(self): def watch(self) -> None:
def process_callback(path: Path): def process_callback(path: Path) -> None:
self._process_single(path) self.process_single(path)
watcher = Watcher(process_callback) watcher = Watcher(process_callback)
watcher.start() watcher.start()
try: try:
while True: while True:
time.sleep(1) time.sleep(1)
except KeyboardInterrupt: except KeyboardInterrupt:
watcher.stop() watcher.stop()
def get_stats(self) -> dict[str, Any]: def get_stats(self, days: int = 30) -> dict[str, Any]:
db_stats = self.db.get_stats() return self.db.get_stats(days)
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,
}

View File

@@ -2,7 +2,7 @@ from __future__ import annotations
from .chunk import SemanticChunker, TextChunk, chunk_text from .chunk import SemanticChunker, TextChunk, chunk_text
from .clean import CleanResult, clean_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__ = [ __all__ = [
"clean_text", "clean_text",
@@ -13,5 +13,5 @@ __all__ = [
"extract_entities", "extract_entities",
"summarize_entities", "summarize_entities",
"Entity", "Entity",
"SciSpacyNER", "ScientificEntityRecognizer",
] ]

View File

@@ -3,7 +3,7 @@ from __future__ import annotations
from dataclasses import dataclass from dataclasses import dataclass
from typing import Any 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.config import settings
from ocr_pipeline.utils.logging import get_logger from ocr_pipeline.utils.logging import get_logger

View File

@@ -24,140 +24,54 @@ def fix_encoding(text: str) -> str:
return text return text
def normalize_unicode(text: str) -> str:
return unicodedata.normalize("NFC", text)
def fix_ligatures(text: str) -> str:
ligatures = {
"": "fi",
"": "fl",
"": "ff",
"": "ffi",
"": "ffl",
"": "ft",
"": "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: def normalize_units(text: str) -> str:
unit_fixes = { number = r"(\d+(?:\.\d+)?)"
r"(\d)\s*u[mM]\b": r"\1 µm", replacements = [
r"(\d)\s*[µu]g\b": r"\1 µg", (rf"{number}\s*[uµ]M\b", r"\1 µM"),
r"(\d)\s*[mn]g\b": r"\1 mg", (rf"{number}\s*[uµ]m\b", r"\1 µm"),
r"(\d)\s*[kK]b\b": r"\1 kb", (rf"{number}\s*[uµ]g\b", r"\1 µg"),
r"(\d)\s*[mM]b\b": r"\1 Mb", (rf"{number}\s*ng\b", r"\1 ng"),
r"(\d)\s*[gG]b\b": r"\1 Gb", (rf"{number}\s*mg\b", r"\1 mg"),
r"(\d)\s*[pP]?[mM]\b": r"\1 pm", (rf"{number}\s*[uµ]L\b", r"\1 µL"),
r"(\d)\s*[nN][mM]\b": r"\1 nm", (rf"{number}\s*mL\b", r"\1 mL"),
r"(\d)\s*[cC][mM]\b": r"\1 cm", (rf"{number}\s*pM\b", r"\1 pM"),
r"(\d)\s*[mM][lL]\b": r"\1 mL", (rf"{number}\s*nM\b", r"\1 nM"),
r"(\d)\s*[uU][lL]\b": r"\1 µL", (rf"{number}\s*mM\b", r"\1 mM"),
r"(\d)\s*°\s*[cC]\b": r"\1 °C", (rf"{number}\s*[kK]b\b", r"\1 kb"),
} (rf"{number}\s*[mM]b\b", r"\1 Mb"),
for pattern, repl in unit_fixes.items(): (rf"{number}\s*[gG]b\b", r"\1 Gb"),
text = re.sub(pattern, repl, text) (rf"{number}\s*°\s*[cC]\b", r"\1 °C"),
return text ]
for pattern, replacement in replacements:
text = re.sub(pattern, replacement, 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)
return text return text
def clean_text(text: str) -> CleanResult: def clean_text(text: str) -> CleanResult:
original_len = len(text) original_length = len(text)
operations = [] operations: list[str] = []
try:
text = fix_encoding(text) text = fix_encoding(text)
finally:
operations.append("fix_encoding") operations.append("fix_encoding")
text = unicodedata.normalize("NFC", text)
text = normalize_unicode(text)
operations.append("normalize_unicode") operations.append("normalize_unicode")
for source, target in {"": "fi", "": "fl", "": "ff", "": "ffi", "": "ffl"}.items():
text = fix_ligatures(text) text = text.replace(source, target)
operations.append("fix_ligatures") operations.append("fix_ligatures")
text = "".join(char for char in text if unicodedata.category(char)[0] != "C" or char in "\n\t")
text = remove_control_chars(text)
operations.append("remove_control_chars") operations.append("remove_control_chars")
text = re.sub(r"(\w+)-\n(\w+)", r"\1\2", text)
text = fix_hyphenation(text)
operations.append("fix_hyphenation") operations.append("fix_hyphenation")
text = re.sub(r"(?m)^\s*[|┃║─━┌┐└┘├┤┬┴┼]+\s*$", "", text)
text = normalize_whitespace(text) text = re.sub(r"(.)\1{10,}", "", text)
operations.append("normalize_whitespace")
text = remove_ocr_artifacts(text)
operations.append("remove_ocr_artifacts") operations.append("remove_ocr_artifacts")
text = re.sub(r"(\d)\s*[×xX]\s*10\s*\^?\s*([+-]?\d+)", r"\1×10^\2", text)
text = fix_scientific_notation(text)
operations.append("fix_scientific_notation") operations.append("fix_scientific_notation")
text = normalize_units(text) text = normalize_units(text)
operations.append("normalize_units") operations.append("normalize_units")
text = re.sub(r"[ \t]+", " ", text)
text = fix_line_breaks(text) text = re.sub(r"\n{3,}", "\n\n", text).strip()
operations.append("fix_line_breaks") operations.append("normalize_whitespace")
logger.debug("text_cleaned", original=original_length, cleaned=len(text), ops=operations)
cleaned_len = len(text) return CleanResult(text=text, original_length=original_length, cleaned_length=len(text), operations=operations)
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,
)

View File

@@ -27,12 +27,14 @@ class EntityResult:
text: str text: str
total_entities: int total_entities: int
entities_by_type: dict[str, int] entities_by_type: dict[str, int]
backend: str = "unknown"
class ScientificEntityRecognizer: class ScientificEntityRecognizer:
def __init__(self): def __init__(self):
self._nlp = None self._nlp = None
self._initialized = False self._initialized = False
self.backend = "unavailable"
def _init(self): def _init(self):
if self._initialized: if self._initialized:
@@ -42,6 +44,7 @@ class ScientificEntityRecognizer:
self._nlp = spacy.load(settings.entities.model) self._nlp = spacy.load(settings.entities.model)
self._nlp.max_length = 2_000_000 self._nlp.max_length = 2_000_000
self._initialized = True self._initialized = True
self.backend = "scispacy"
logger.info("scispacy_model_loaded", model=settings.entities.model) logger.info("scispacy_model_loaded", model=settings.entities.model)
except Exception as e: except Exception as e:
logger.warning("scispacy_load_failed", error=str(e), model=settings.entities.model) logger.warning("scispacy_load_failed", error=str(e), model=settings.entities.model)
@@ -52,6 +55,7 @@ class ScientificEntityRecognizer:
try: try:
self._nlp = spacy.blank("en") self._nlp = spacy.blank("en")
self._nlp.add_pipe("sentencizer") self._nlp.add_pipe("sentencizer")
self.backend = "blank_spacy"
logger.info("using_fallback_blank_model") logger.info("using_fallback_blank_model")
except Exception: except Exception:
self._nlp = None self._nlp = None
@@ -99,6 +103,7 @@ class ScientificEntityRecognizer:
text=text, text=text,
total_entities=len(entities), total_entities=len(entities),
entities_by_type=by_type, entities_by_type=by_type,
backend=self.backend,
) )
except Exception as e: except Exception as e:
logger.warning("entity_extraction_failed", error=str(e)) logger.warning("entity_extraction_failed", error=str(e))
@@ -141,15 +146,14 @@ class ScientificEntityRecognizer:
class RegexEntityRecognizer: class RegexEntityRecognizer:
PATTERNS = { PATTERNS = {
"GENE": [ "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)\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",
], ],
"PROTEIN": [ "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", 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": [ "CHEMICAL": [
r"\b(?:DMSO|PBS|EDTA|Tris|HEPES|SDS|DTT|BME|NaCl|KCl|MgCl2|CaCl2|NaOH|HCl|H2SO4|HNO3|EtOH|MeOH|IPA|DMSO|DMF|DMA|THF|DCM|CHCl3|CH2Cl2|EtOAc|hexane|pentane|acetone|acetonitrile|water|H2O|buffer|media|serum|FBS|BSA|pen/strep|penicillin|streptomycin|trypsin|EDTA|collagenase|dispase|accutase)\b", r"\b(?:DMSO|PBS|EDTA|Tris|HEPES|SDS|DTT|BME|NaCl|KCl|MgCl2|CaCl2|NaOH|HCl|H2SO4|HNO3|EtOH|MeOH|IPA|DMSO|DMF|DMA|THF|DCM|CHCl3|CH2Cl2|EtOAc|hexane|pentane|acetone|acetonitrile|water|H2O|buffer|media|serum|FBS|BSA|pen/strep|penicillin|streptomycin|trypsin|EDTA|collagenase|dispase|accutase)\b",
r"\b(?:[A-Z][a-z]?\d*(?:\([^)]+\))?(?:\s*[+\-]\s*[A-Z][a-z]?\d*(?:\([^)]+\))?)*)\b",
], ],
"CONCENTRATION": [ "CONCENTRATION": [
r"\b\d+(?:\.\d+)?\s*(?:[µumMpnc]?[Mm]|[µumMpnc]?[Mm]/[Ll]|[µumMpnc]?[Mm]\s*[Ll]?|[µumMpnc]?[gG]\s*/\s*[Ll]|[µumMpnc]?[gG]\s*/\s*[mM][lL]|[µumMpnc]?[gG]\s*/\s*[dL]|[µumMpnc]?[gG]\s*/\s*[mM][lL]|[µumMpnc]?[gG]\s*/\s*[dD][lL]|[µumMpnc]?[gG]\s*/\s*100\s*[mM][lL]|[µumMpnc]?[gG]\s*/\s*[kK][gG])\b", r"\b\d+(?:\.\d+)?\s*(?:[µumMpnc]?[Mm]|[µumMpnc]?[Mm]/[Ll]|[µumMpnc]?[Mm]\s*[Ll]?|[µumMpnc]?[gG]\s*/\s*[Ll]|[µumMpnc]?[gG]\s*/\s*[mM][lL]|[µumMpnc]?[gG]\s*/\s*[dL]|[µumMpnc]?[gG]\s*/\s*[mM][lL]|[µumMpnc]?[gG]\s*/\s*[dD][lL]|[µumMpnc]?[gG]\s*/\s*100\s*[mM][lL]|[µumMpnc]?[gG]\s*/\s*[kK][gG])\b",
@@ -177,9 +181,14 @@ class RegexEntityRecognizer:
return entities return entities
_scientific_recognizer = ScientificEntityRecognizer()
def extract_entities(text: str) -> EntityResult: def extract_entities(text: str) -> EntityResult:
recognizer = ScientificEntityRecognizer() if not settings.entities.enabled:
result = recognizer.extract(text) return EntityResult(entities=[], text=text, total_entities=0, entities_by_type={})
result = _scientific_recognizer.extract(text)
regex_entities = RegexEntityRecognizer.extract(text) regex_entities = RegexEntityRecognizer.extract(text)
all_entities = result.entities + regex_entities all_entities = result.entities + regex_entities
@@ -194,6 +203,7 @@ def extract_entities(text: str) -> EntityResult:
text=text, text=text,
total_entities=len(all_entities), total_entities=len(all_entities),
entities_by_type=by_type, 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)} |") lines.append(f"| {ents[0].text} | {label} | {len(ents)} | {'; '.join(contexts)} |")
return "\n".join(lines) 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)

View File

@@ -2,16 +2,14 @@ from __future__ import annotations
from .db import Database, ProcessedFile, get_db from .db import Database, ProcessedFile, get_db
from .logging import ProgressTracker, create_progress, get_logger, setup_logging from .logging import ProgressTracker, create_progress, get_logger, setup_logging
from .migrate import compute_file_hash, migrate_notebook, parse_notebook_for_images
__all__ = [ __all__ = [
"setup_logging", "setup_logging",
"get_logger", "get_logger",
"create_progress", "create_progress",
"ProgressTracker", "ProgressTracker",
"parse_notebook_for_images",
"migrate_notebook",
"compute_file_hash",
"get_db", "get_db",
"Database", "Database",
"ProcessedFile", "ProcessedFile",

View File

@@ -2,6 +2,7 @@ from __future__ import annotations
import sqlite3 import sqlite3
import time import time
from collections.abc import Iterator
from contextlib import contextmanager from contextlib import contextmanager
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
@@ -12,32 +13,34 @@ from ocr_pipeline.utils.logging import get_logger
logger = get_logger(__name__) logger = get_logger(__name__)
@dataclass @dataclass(frozen=True)
class ProcessedFile: class ProcessedFile:
"""Legacy value object retained for import compatibility."""
path: str path: str
hash: str hash: str
timestamp: float timestamp: float
run_id: int run_id: int | None = None
size: int size: int = 0
mtime: float mtime: float = 0.0
class Database: class Database:
def __init__(self, db_path: Path): def __init__(self, db_path: Path) -> None:
self.db_path = Path(db_path).expanduser().resolve() self.db_path = db_path.expanduser().resolve()
self.db_path.parent.mkdir(parents=True, exist_ok=True) self.db_path.parent.mkdir(parents=True, exist_ok=True)
self._init_db() self._init_db()
@contextmanager @contextmanager
def _conn(self): def _conn(self) -> Iterator[sqlite3.Connection]:
conn = sqlite3.connect(self.db_path) conn = sqlite3.connect(self.db_path)
conn.row_factory = sqlite3.Row conn.row_factory = sqlite3.Row
try: try:
yield conn yield conn
conn.commit()
finally: finally:
conn.close() conn.close()
def _init_db(self): def _init_db(self) -> None:
with self._conn() as conn: with self._conn() as conn:
conn.executescript(""" conn.executescript("""
CREATE TABLE IF NOT EXISTS processed_files ( CREATE TABLE IF NOT EXISTS processed_files (
@@ -47,9 +50,9 @@ class Database:
ocr_engine TEXT, ocr_engine TEXT,
confidence REAL, confidence REAL,
chunks_created INTEGER DEFAULT 0, chunks_created INTEGER DEFAULT 0,
run_id INTEGER,
created_at REAL DEFAULT (strftime('%s', 'now')) created_at REAL DEFAULT (strftime('%s', 'now'))
); );
CREATE TABLE IF NOT EXISTS runs ( CREATE TABLE IF NOT EXISTS runs (
id INTEGER PRIMARY KEY AUTOINCREMENT, id INTEGER PRIMARY KEY AUTOINCREMENT,
date TEXT NOT NULL, date TEXT NOT NULL,
@@ -62,96 +65,49 @@ class Database:
completed_at REAL, completed_at REAL,
UNIQUE(date, run_number) UNIQUE(date, run_number)
); );
CREATE INDEX IF NOT EXISTS idx_processed_hash ON processed_files(hash); CREATE INDEX IF NOT EXISTS idx_processed_hash ON processed_files(hash);
CREATE INDEX IF NOT EXISTS idx_runs_date ON runs(date); 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: 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: with self._conn() as conn:
if path: if path is None:
cursor = conn.execute( row = conn.execute("SELECT 1 FROM processed_files WHERE hash = ?", (file_hash,)).fetchone()
"SELECT 1 FROM processed_files WHERE hash = ? AND path = ?",
(file_hash, str(path)),
)
else: else:
cursor = conn.execute( row = conn.execute("SELECT 1 FROM processed_files WHERE hash = ? AND path = ?", (file_hash, str(path))).fetchone()
"SELECT 1 FROM processed_files WHERE hash = ?", return row is not None
(file_hash,),
)
return cursor.fetchone() is not None
def get_processed(self, file_hash: str) -> dict[str, Any] | None: def get_processed(self, file_hash: str) -> dict[str, Any] | None:
with self._conn() as conn: with self._conn() as conn:
cursor = conn.execute("SELECT * FROM processed_files WHERE hash = ?", (file_hash,)) row = conn.execute("SELECT * FROM processed_files WHERE hash = ?", (file_hash,)).fetchone()
row = cursor.fetchone()
return dict(row) if row else None return dict(row) if row else None
def mark_processed( 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:
self,
file_hash: str,
path: str,
timestamp: float | str,
engine: str = "",
confidence: float = 0.0,
chunks: int = 0,
run_id: int = 0,
):
with self._conn() as conn: with self._conn() as conn:
conn.execute( conn.execute("""INSERT OR REPLACE INTO processed_files
"""INSERT OR REPLACE INTO processed_files (hash, path, timestamp, ocr_engine, confidence, chunks_created, run_id)
(hash, path, timestamp, ocr_engine, confidence, chunks_created) VALUES (?, ?, ?, ?, ?, ?, ?)""", (file_hash, path, str(timestamp), engine, confidence, chunks, run_id))
VALUES (?, ?, ?, ?, ?, ?)""",
(file_hash, path, str(timestamp), engine, confidence, chunks), 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]: def get_stats(self, days: int = 30) -> dict[str, Any]:
cutoff = time.time() - days * 86400
with self._conn() as conn: with self._conn() as conn:
cursor = conn.execute(""" 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()
SELECT runs = conn.execute("SELECT COUNT(*) AS total_runs, MAX(completed_at) AS latest_run FROM runs WHERE started_at >= ?", (cutoff,)).fetchone()
COUNT(*) as total_processed, return {**dict(files), **dict(runs)}
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,
),
)
_db_instance: Database | None = None _db_instance: Database | None = None
@@ -161,6 +117,5 @@ def get_db() -> Database:
global _db_instance global _db_instance
if _db_instance is None: if _db_instance is None:
from ocr_pipeline.config import settings from ocr_pipeline.config import settings
_db_instance = Database(settings.watch_db_path) _db_instance = Database(settings.watch_db_path)
return _db_instance return _db_instance

View File

@@ -7,9 +7,7 @@ from pathlib import Path
from typing import Any from typing import Any
import nbformat import nbformat
import platformdirs
from ocr_pipeline.config import settings
from ocr_pipeline.utils.logging import get_logger from ocr_pipeline.utils.logging import get_logger
logger = get_logger(__name__) logger = get_logger(__name__)
@@ -104,8 +102,11 @@ def migrate_notebook(notebook_path: Path, pipeline: Any) -> dict[str, Any]:
continue continue
logger.info("migrating_image", path=str(img_path)) 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 results["processed"] += 1
else:
results["skipped"] += 1
except Exception as e: except Exception as e:
logger.error("migration_failed", path=str(img_path), error=str(e)) logger.error("migration_failed", path=str(img_path), error=str(e))

View File

@@ -1,144 +1,5 @@
from __future__ import annotations """Compatibility module for the canonical watcher package."""
import threading from .watch.watcher import ScreenshotHandler, Watcher
import time
from collections.abc import Callable
from pathlib import Path
from watchdog.events import FileSystemEventHandler __all__ = ["ScreenshotHandler", "Watcher"]
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")

View File

@@ -74,7 +74,7 @@ class Watcher:
if self.running: if self.running:
return return
self.run_id = self.db.start_run() self.run_id, _ = self.db.start_run()
self.running = True self.running = True
self.handler = ScreenshotHandler(self._process_file) self.handler = ScreenshotHandler(self._process_file)
@@ -112,29 +112,14 @@ class Watcher:
logger.error("watch_process_failed", path=str(path), error=str(e)) logger.error("watch_process_failed", path=str(path), error=str(e))
failed += 1 failed += 1
self.db.complete_run( if self.run_id is not None:
time.strftime("%Y-%m-%d"), self.db.complete_run(self.run_id, {"total": processed + failed, "successful": processed, "failed": failed, "chunks": 0})
self.run_id,
{"total": processed, "successful": processed, "failed": failed, "chunks": 0},
)
logger.info("watcher_stopped", run_id=self.run_id, processed=processed, failed=failed) logger.info("watcher_stopped", run_id=self.run_id, processed=processed, failed=failed)
def _process_file(self, path: Path) -> None: def _process_file(self, path: Path) -> None:
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)) logger.info("processing_new_file", path=str(path))
# The pipeline owns deduplication and database writes, preventing double marks.
self.process_callback(path) 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: def stop(self) -> None:
self.running = False self.running = False

29
tests/test_output.py Normal file
View 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

6120
uv.lock generated

File diff suppressed because it is too large Load Diff