OCR-to-RAG pipeline for life science screenshots
PaddleOCR with preprocessing, scispaCy NER, figure/table detection, citation extraction, and chunked Markdown output with frontmatter. Includes watch mode and notebook reprocessing.
This commit is contained in:
246
src/ocr_pipeline/config.py
Normal file
246
src/ocr_pipeline/config.py
Normal file
@@ -0,0 +1,246 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any, ClassVar
|
||||
|
||||
import yaml
|
||||
from pydantic import Field, field_validator
|
||||
from pydantic_settings import BaseSettings, SettingsConfigDict
|
||||
|
||||
|
||||
class InputConfig(BaseSettings):
|
||||
paths: list[str] = Field(default_factory=lambda: ["~/Pictures"])
|
||||
patterns: list[str] = Field(
|
||||
default_factory=lambda: ["SCR-*.png", "*.jpg", "*.jpeg", "*.tiff", "*.bmp"]
|
||||
)
|
||||
recursive: bool = True
|
||||
|
||||
@field_validator("paths", mode="before")
|
||||
@classmethod
|
||||
def expand_paths(cls, v: list[str]) -> list[str]:
|
||||
return [os.path.expanduser(p) for p in v]
|
||||
|
||||
|
||||
class PreprocessConfig(BaseSettings):
|
||||
deskew: bool = True
|
||||
denoise: bool = True
|
||||
clahe: bool = True
|
||||
adaptive_threshold: bool = True
|
||||
remove_lines: bool = True
|
||||
max_dimension: int = 4096
|
||||
|
||||
|
||||
class OCRConfig(BaseSettings):
|
||||
engine: str = "paddleocr" # paddleocr | tesseract | auto
|
||||
languages: list[str] = Field(default_factory=lambda: ["en", "latin"])
|
||||
use_gpu: bool = False
|
||||
use_angle_cls: bool = True
|
||||
det_db_thresh: float = 0.3
|
||||
det_db_box_thresh: float = 0.6
|
||||
det_db_unclip_ratio: float = 1.5
|
||||
rec_batch_num: int = 6
|
||||
cpu_threads: int = 4
|
||||
preprocess: PreprocessConfig = Field(default_factory=PreprocessConfig)
|
||||
|
||||
|
||||
class ProcessingConfig(BaseSettings):
|
||||
workers: int = 4
|
||||
batch_size: int = 10
|
||||
retry_attempts: int = 2
|
||||
timeout_per_image: int = 120
|
||||
skip_existing: bool = True
|
||||
|
||||
|
||||
class DetectorConfig(BaseSettings):
|
||||
enabled: bool = True
|
||||
model: str = ""
|
||||
confidence_threshold: float = 0.7
|
||||
|
||||
|
||||
class FigureDetectorConfig(DetectorConfig):
|
||||
model: str = "lp://PubLayNet/faster_rcnn_R_50_FPN_3x/config"
|
||||
confidence_threshold: float = 0.7
|
||||
|
||||
|
||||
class TableDetectorConfig(DetectorConfig):
|
||||
model: str = "microsoft/table-transformer-detection"
|
||||
confidence_threshold: float = 0.7
|
||||
|
||||
|
||||
class CaptionDetectorConfig(DetectorConfig):
|
||||
model: str = "lp://PubLayNet/faster_rcnn_R_50_FPN_3x/config"
|
||||
confidence_threshold: float = 0.5
|
||||
|
||||
|
||||
class DetectorsConfig(BaseSettings):
|
||||
figures: FigureDetectorConfig = Field(default_factory=FigureDetectorConfig)
|
||||
tables: TableDetectorConfig = Field(default_factory=TableDetectorConfig)
|
||||
captions: CaptionDetectorConfig = Field(default_factory=CaptionDetectorConfig)
|
||||
|
||||
|
||||
class EntitiesConfig(BaseSettings):
|
||||
enabled: bool = True
|
||||
model: str = "en_core_sci_lg"
|
||||
types: list[str] = Field(
|
||||
default_factory=lambda: [
|
||||
"GENE",
|
||||
"PROTEIN",
|
||||
"CHEMICAL",
|
||||
"SPECIES",
|
||||
"DISEASE",
|
||||
"CELL_LINE",
|
||||
"ORGANISM",
|
||||
"CELL_TYPE",
|
||||
]
|
||||
)
|
||||
merge_entities: bool = True
|
||||
|
||||
|
||||
class CitationsConfig(BaseSettings):
|
||||
regex_enabled: bool = True
|
||||
grobid_enabled: bool = False
|
||||
grobid_url: str = "http://localhost:8070"
|
||||
grobid_timeout: int = 30
|
||||
|
||||
|
||||
class ChunkingConfig(BaseSettings):
|
||||
chunk_size: int = 1000
|
||||
chunk_overlap: int = 200
|
||||
separators: list[str] = Field(default_factory=lambda: ["\n\n", "\n", ". ", " ", ""])
|
||||
keep_separator: bool = True
|
||||
|
||||
|
||||
class FrontmatterConfig(BaseSettings):
|
||||
fields: list[str] = Field(
|
||||
default_factory=lambda: [
|
||||
"source_path",
|
||||
"source_hash",
|
||||
"timestamp",
|
||||
"ocr_engine",
|
||||
"ocr_confidence_mean",
|
||||
"language",
|
||||
"detected_entities",
|
||||
"has_figures",
|
||||
"has_tables",
|
||||
"citations_found",
|
||||
"chunk_index",
|
||||
"total_chunks",
|
||||
]
|
||||
)
|
||||
include_raw_text: bool = False
|
||||
|
||||
|
||||
class OutputConfig(BaseSettings):
|
||||
base_directory: str = "./data/ocr_output"
|
||||
format: str = "markdown"
|
||||
organize_by: str = "date_run" # date_run | source_dir | flat
|
||||
frontmatter: FrontmatterConfig = Field(default_factory=FrontmatterConfig)
|
||||
|
||||
|
||||
class WatchConfig(BaseSettings):
|
||||
enabled: bool = True
|
||||
debounce_seconds: int = 5
|
||||
ignore_patterns: list[str] = Field(
|
||||
default_factory=lambda: [
|
||||
".DS_Store",
|
||||
"*.tmp",
|
||||
"*.partial",
|
||||
"*.crdownload",
|
||||
]
|
||||
)
|
||||
db_path: str = "./data/processed_files.db"
|
||||
poll_interval: float = 1.0
|
||||
|
||||
|
||||
class _BaseSettings(BaseSettings):
|
||||
model_config = SettingsConfigDict(
|
||||
env_file=".env",
|
||||
env_nested_delimiter="__",
|
||||
extra="ignore",
|
||||
)
|
||||
|
||||
|
||||
def _load_yaml_config(path: str = "config.yaml") -> dict[str, Any]:
|
||||
"""Load YAML config file if it exists."""
|
||||
config_path = Path(path)
|
||||
if not config_path.is_absolute():
|
||||
config_path = Path.cwd() / config_path
|
||||
if config_path.exists():
|
||||
with open(config_path) as f:
|
||||
return yaml.safe_load(f) or {}
|
||||
return {}
|
||||
|
||||
|
||||
def _build_settings_from_yaml(yaml_dict: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Convert YAML dict to properly typed nested models for Settings."""
|
||||
if not yaml_dict:
|
||||
return {}
|
||||
|
||||
result = {}
|
||||
|
||||
# Map YAML keys to model classes
|
||||
model_map = {
|
||||
"input": InputConfig,
|
||||
"ocr": OCRConfig,
|
||||
"processing": ProcessingConfig,
|
||||
"detectors": DetectorsConfig,
|
||||
"entities": EntitiesConfig,
|
||||
"citations": CitationsConfig,
|
||||
"chunking": ChunkingConfig,
|
||||
"output": OutputConfig,
|
||||
"watch": WatchConfig,
|
||||
}
|
||||
|
||||
for key, model_class in model_map.items():
|
||||
if key in yaml_dict:
|
||||
data = yaml_dict[key]
|
||||
# Special handling for output config to transform frontmatter list -> dict
|
||||
if key == "output" and "frontmatter" in data and isinstance(data["frontmatter"], list):
|
||||
data = {**data, "frontmatter": {"fields": data["frontmatter"]}}
|
||||
result[key] = model_class(**data)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class Settings(_BaseSettings):
|
||||
input: InputConfig = Field(default_factory=InputConfig)
|
||||
ocr: OCRConfig = Field(default_factory=OCRConfig)
|
||||
processing: ProcessingConfig = Field(default_factory=ProcessingConfig)
|
||||
detectors: DetectorsConfig = Field(default_factory=DetectorsConfig)
|
||||
entities: EntitiesConfig = Field(default_factory=EntitiesConfig)
|
||||
citations: CitationsConfig = Field(default_factory=CitationsConfig)
|
||||
chunking: ChunkingConfig = Field(default_factory=ChunkingConfig)
|
||||
output: OutputConfig = Field(default_factory=OutputConfig)
|
||||
watch: WatchConfig = Field(default_factory=WatchConfig)
|
||||
|
||||
# Store YAML config for manual override
|
||||
_yaml_config: ClassVar[dict[str, Any]] = {}
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
if not Settings._yaml_config:
|
||||
Settings._yaml_config = _load_yaml_config()
|
||||
|
||||
# Build nested models from YAML
|
||||
yaml_models = _build_settings_from_yaml(Settings._yaml_config)
|
||||
super().__init__(**yaml_models, **kwargs)
|
||||
|
||||
def update_from_yaml(self, path: str) -> None:
|
||||
"""Reload settings from a YAML file."""
|
||||
Settings._yaml_config = _load_yaml_config(path)
|
||||
# Re-initialize with new config
|
||||
yaml_models = _build_settings_from_yaml(Settings._yaml_config)
|
||||
for key, value in yaml_models.items():
|
||||
if hasattr(self, key):
|
||||
setattr(self, key, value)
|
||||
|
||||
@property
|
||||
def output_dir(self) -> Path:
|
||||
return Path(self.output.base_directory).expanduser().resolve()
|
||||
|
||||
@property
|
||||
def watch_db_path(self) -> Path:
|
||||
return Path(self.watch.db_path).expanduser().resolve()
|
||||
|
||||
|
||||
settings = Settings()
|
||||
Reference in New Issue
Block a user