170 lines
5.0 KiB
Python
170 lines
5.0 KiB
Python
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()
|