batch/dataset_processor.py

"""Batch dataset processing and reproducibility manifest export."""

from __future__ import annotations

import json
import shutil
import tempfile
import zipfile
from pathlib import Path
from typing import Any, Iterable

import cv2
import numpy as np
from PIL import Image

from cv_ops.analysis import stats
from filters.registry import apply_definition, load_definition


IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"}


def _iter_input_files(files: list[str] | None, directory: str | None, workdir: Path) -> Iterable[Path]:
    if directory:
        root = Path(directory).expanduser()
        if root.exists():
            yield from (p for p in root.rglob("*") if p.suffix.lower() in IMAGE_SUFFIXES)
    for file in files or []:
        path = Path(file)
        if path.suffix.lower() == ".zip":
            with zipfile.ZipFile(path) as zf:
                zf.extractall(workdir / path.stem)
            yield from (p for p in (workdir / path.stem).rglob("*") if p.suffix.lower() in IMAGE_SUFFIXES)
        elif path.suffix.lower() in IMAGE_SUFFIXES:
            yield path


def process_dataset(files: list[str] | None, directory: str | None, filter_name: str, progress: Any = None) -> str:
    definition = load_definition(filter_name)
    temp_root = Path(tempfile.mkdtemp(prefix="cv_lab_batch_"))
    out_dir = temp_root / "processed"
    out_dir.mkdir()
    manifest: list[dict[str, Any]] = []
    inputs = list(_iter_input_files(files, directory, temp_root))
    total = max(1, len(inputs))
    for idx, path in enumerate(inputs):
        if progress:
            progress((idx + 1) / total, desc=f"Processing {path.name}")
        record: dict[str, Any] = {"filename": path.name, "filter": filter_name}
        try:
            img = np.array(Image.open(path).convert("RGB"))
            result = apply_definition(img, definition)
            out_path = out_dir / f"{path.stem}_processed.png"
            cv2.imwrite(str(out_path), cv2.cvtColor(result, cv2.COLOR_RGB2BGR))
            record.update({"status": "ok", "output": out_path.name, "stats": stats(result)})
        except Exception as exc:
            record.update({"status": "error", "error": str(exc)})
        manifest.append(record)
    (out_dir / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
    archive = shutil.make_archive(str(temp_root / "cv_lab_processed"), "zip", out_dir)
    return archive