272 lines
8.1 KiB
Python
272 lines
8.1 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import platform
|
|
from pathlib import Path
|
|
from typing import Any, ClassVar
|
|
|
|
import platformdirs
|
|
import yaml
|
|
from pydantic import Field, field_validator
|
|
from pydantic_settings import BaseSettings, SettingsConfigDict
|
|
|
|
|
|
APP_NAME = "ocr-pipeline"
|
|
APP_AUTHOR = "aman"
|
|
|
|
|
|
def _get_platform_ignore_patterns() -> list[str]:
|
|
"""Get platform-specific ignore patterns for file watching."""
|
|
patterns = [
|
|
"*.tmp",
|
|
"*.partial",
|
|
"*.crdownload",
|
|
"*.part",
|
|
"~$*", # MS Office temp files
|
|
"*.swp",
|
|
"*.swo", # Vim swap files
|
|
]
|
|
system = platform.system().lower()
|
|
if system == "darwin":
|
|
patterns.extend([".DS_Store", "._*", ".Spotlight-V100", ".Trashes", ".fseventsd"])
|
|
elif system == "windows":
|
|
patterns.extend(["Thumbs.db", "desktop.ini", "*.lnk", "$RECYCLE.BIN"])
|
|
else:
|
|
# Linux/Unix
|
|
patterns.extend([".directory", "*.swp", "*.swo", ".git"])
|
|
return patterns
|
|
|
|
|
|
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 = Field(
|
|
default_factory=lambda: str(platformdirs.user_data_dir(APP_NAME, APP_AUTHOR))
|
|
)
|
|
format: str = "markdown"
|
|
organize_by: str = "date_run" # date_run | source_dir | flat
|
|
frontmatter: FrontmatterConfig = Field(default_factory=FrontmatterConfig)
|
|
|
|
|
|
class WatchConfig(BaseSettings):
|
|
enabled: bool = True
|
|
debounce_seconds: int = 5
|
|
ignore_patterns: list[str] = Field(default_factory=_get_platform_ignore_patterns)
|
|
db_path: str = Field(
|
|
default_factory=lambda: str(Path(platformdirs.user_data_dir(APP_NAME, APP_AUTHOR)) / "processed_files.db")
|
|
)
|
|
poll_interval: float = 1.0
|
|
|
|
|
|
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()
|