from __future__ import annotations import sys from pathlib import Path import typer from rich.console import Console from rich.table import Table from ocr_pipeline.config import settings from ocr_pipeline.pipeline import OCRPipeline from ocr_pipeline.utils.logging import get_logger, setup_logging from ocr_pipeline.utils.migrate import migrate_notebook app = typer.Typer( name="ocr-pipeline", help="OCR pipeline for screenshots -> RAG-ready Markdown", add_completion=False, ) console = Console() logger = get_logger(__name__) @app.callback() def callback( config: Path | None = typer.Option(None, "--config", "-c", help="Config file path"), log_level: str = typer.Option("INFO", "--log-level", "-l", help="Log level"), log_file: Path | None = typer.Option(None, "--log-file", help="Log file path"), ): """OCR Pipeline - Extract text from screenshots for RAG""" setup_logging(log_level, log_file) if config: settings.update_from_yaml(config) @app.command() def run( input_dir: Path | None = typer.Option(None, "--input-dir", "-i", help="Input directory"), output_dir: Path | None = typer.Option(None, "--output-dir", "-o", help="Output directory"), workers: int | None = typer.Option( None, "--workers", "-w", help="Number of worker processes" ), engine: str | None = typer.Option( None, "--engine", "-e", help="OCR engine (paddleocr/tesseract/auto)" ), exclude: list[str] = typer.Option([], "--exclude", help="Glob pattern to exclude; repeatable"), force: bool = typer.Option(False, "--force", help="Reprocess files already recorded as processed"), ): """Run OCR pipeline on all screenshots""" if input_dir: settings.input.paths = [str(input_dir)] if output_dir: settings.output.base_directory = str(output_dir) if workers: settings.processing.workers = workers if engine: settings.ocr.engine = engine if exclude: settings.input.exclude_patterns = exclude if force: settings.processing.skip_existing = False pipeline = OCRPipeline() result = pipeline.process_all() table = Table(title="Pipeline Results") table.add_column("Metric", style="cyan") table.add_column("Value", style="green") table.add_row("Total Images", str(result.total_images)) table.add_row("Successful", str(result.successful)) table.add_row("Failed", str(result.failed)) table.add_row("Chunks Created", str(result.chunks_created)) table.add_row("Output Files", str(len(result.output_files))) console.print(table) if result.errors: console.print("\n[red]Errors:[/red]") for err in result.errors: console.print(f" {err['path']}: {err['error']}") if result.failed > 0: sys.exit(1) @app.command() def watch(): """Watch for new screenshots and process automatically""" pipeline = OCRPipeline() console.print("[green]Starting watch mode...[/green]") console.print("Press Ctrl+C to stop") try: pipeline.watch() except KeyboardInterrupt: console.print("\n[yellow]Stopping watcher...[/yellow]") @app.command() def reprocess( notebook: Path = typer.Argument(..., help="Path to Jupyter notebook"), ): """Reprocess screenshots from a Jupyter notebook""" pipeline = OCRPipeline() console.print(f"[green]Migrating from notebook: {notebook}[/green]") result = migrate_notebook(notebook, pipeline) table = Table(title="Migration Results") table.add_column("Metric", style="cyan") table.add_column("Value", style="green") for key, value in result.items(): if key != "errors": table.add_row(key.capitalize(), str(value)) console.print(table) if result.get("errors"): console.print("\n[red]Errors:[/red]") for err in result["errors"]: console.print(f" {err['path']}: {err['error']}") @app.command() def status(): """Show pipeline status and statistics""" pipeline = OCRPipeline() stats = pipeline.get_stats() table = Table(title="Pipeline Status") table.add_column("Metric", style="cyan") table.add_column("Value", style="green") for key, value in stats.items(): table.add_row(key.replace("_", " ").title(), str(value)) console.print(table) @app.command() def stats(days: int = typer.Option(30, "--days", "-d", help="Number of days to analyze")): """Show processing statistics""" pipeline = OCRPipeline() stats = pipeline.get_stats(days) table = Table(title=f"Statistics (last {days} days)") table.add_column("Metric", style="cyan") table.add_column("Value", style="green") for key, value in stats.items(): table.add_row(key.replace("_", " ").title(), str(value)) console.print(table) @app.command() def config(): """Show current configuration""" console.print("[cyan]Current Configuration:[/cyan]") console.print(settings.model_dump_json(indent=2)) def main(): app() if __name__ == "__main__": main()