from __future__ import annotations import hashlib import os import platform from pathlib import Path from typing import Any from ocr_pipeline.utils.logging import get_logger logger = get_logger(__name__) def _get_pictures_dir() -> Path: """Get platform-appropriate Pictures directory.""" system = platform.system().lower() if system == "darwin": return Path.home() / "Pictures" elif system == "windows": # Windows: Pictures folder in user profile return Path.home() / "Pictures" else: # Linux/Unix: XDG Pictures dir or ~/Pictures pictures = Path.home() / "Pictures" if pictures.exists(): return pictures # Fallback to XDG_PICTURES_DIR xdg = os.environ.get("XDG_PICTURES_DIR") if xdg: return Path(xdg) return pictures def parse_notebook_for_images(notebook_path: Path) -> list[Path]: image_paths = [] try: import nbformat with open(notebook_path) as f: nb = nbformat.read(f, as_version=4) for cell in nb.cells: if cell.cell_type == "code": for output in cell.outputs: if output.output_type == "stream" and output.name == "stdout": text = output.text for line in text.split("\n"): line = line.strip() if line.startswith("/") and any( line.endswith(ext) for ext in (".png", ".jpg", ".jpeg", ".tiff", ".bmp") ): path = Path(line) if path.exists(): image_paths.append(path) elif "SCR-" in line and any( line.endswith(ext) for ext in (".png", ".jpg", ".jpeg") ): for part in line.split(): if part.startswith("SCR-") and any( part.endswith(ext) for ext in (".png", ".jpg", ".jpeg") ): path = _get_pictures_dir() / part if path.exists(): image_paths.append(path) except Exception as e: logger.warning("notebook_parse_failed", path=str(notebook_path), error=str(e)) return list(set(image_paths)) def compute_file_hash(path: Path) -> str: hasher = hashlib.sha256() with open(path, "rb") as f: for chunk in iter(lambda: f.read(8192), b""): hasher.update(chunk) return hasher.hexdigest() def migrate_notebook(notebook_path: Path, pipeline: Any) -> dict[str, Any]: logger.info("starting_notebook_migration", notebook=str(notebook_path)) image_paths = parse_notebook_for_images(notebook_path) logger.info("found_images_in_notebook", count=len(image_paths)) results = { "total": len(image_paths), "processed": 0, "skipped": 0, "failed": 0, "errors": [], } for img_path in image_paths: try: img_hash = compute_file_hash(img_path) existing = pipeline.db.get_processed(img_hash) if existing: logger.info("skipping_already_processed", path=str(img_path), hash=img_hash[:16]) results["skipped"] += 1 continue logger.info("migrating_image", path=str(img_path)) output_files = pipeline.process_single(img_path) if output_files: results["processed"] += 1 else: results["skipped"] += 1 except Exception as e: logger.error("migration_failed", path=str(img_path), error=str(e)) results["failed"] += 1 results["errors"].append({"path": str(img_path), "error": str(e)}) logger.info("migration_complete", **results) return results