Aman Nalakath
c5bf5d2b0c
Some checks failed
tests / core (macos-latest, py3.11) (push) Has been cancelled
tests / core (macos-latest, py3.12) (push) Has been cancelled
tests / core (macos-latest, py3.13) (push) Has been cancelled
tests / core (ubuntu-latest, py3.11) (push) Has been cancelled
tests / core (ubuntu-latest, py3.12) (push) Has been cancelled
tests / core (ubuntu-latest, py3.13) (push) Has been cancelled
tests / doctor CLI smoke test (push) Has been cancelled
- Delete notebooks/ (personal scratch artifacts; gitignored) - Move nbformat to --extra migrate; lazy imports in migrate.py and cli.py so base CLI loads without it - Restore langchain-text-splitters in core (chunking dependency)
118 lines
4.0 KiB
Python
118 lines
4.0 KiB
Python
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
|