photo-pipeline/quality_scan.py

"""photo-pipeline: quality scan flow (v3) using CleanVision.

Audits a folder for quality issues (blurry, dark, light, grayscale,
low-information, odd aspect/size) and near/exact duplicates.

Writes a per-image verdict: keep / review / delete-candidate, and moves
files into the /mnt/data/{01_keep,02_review,03_delete} staging dirs.
Nothing is deleted — 03_delete is a holding area for human confirmation.
"""

import shutil
from pathlib import Path

from prefect import flow, task

STAGING = Path("/mnt/data")
KEEP = STAGING / "01_keep"
REVIEW = STAGING / "02_review"
DELETE = STAGING / "03_delete"

# Issue types that warrant deletion-candidate vs review
HARD_ISSUES = {"dark", "light", "low_information", "blurry", "grayscale"}
SOFT_ISSUES = {"odd_aspect_ratio", "odd_size"}


@task
def audit_folder(base_dir: str, issue_types: list[str] | None = None) -> dict:
    """Run CleanVision audit on a folder. Returns issue summary + per-image issues."""
    from cleanvision import Imagelab

    imagelab = Imagelab(data_path=base_dir)
    if issue_types:
        imagelab.find_issues(issue_types=issue_types)
    else:
        imagelab.find_issues()
    summary = imagelab.issue_summary.to_dict("records")

    # imagelab.issues is ONE DataFrame: cols like dark_score/is_dark_issue
    df = imagelab.issues
    per_image = {}
    for idx, row in df.iterrows():
        name = idx
        for col in df.columns:
            if col.startswith("is_") and col.endswith("_issue") and row[col]:
                issue_type = col[len("is_"):-len("_issue")]
                per_image.setdefault(name, []).append(issue_type)
    return {"summary": summary, "per_image": per_image}


@task
def classify_and_sort(base_dir: str, per_image: dict, move: bool = True) -> dict:
    """Classify each image and (optionally) move into staging dirs."""
    root = Path(base_dir)
    images = [p for p in root.rglob("*") if p.is_file()]
    counts = {"keep": 0, "review": 0, "delete_candidate": 0, "skipped": 0}
    decisions = {}

    for p in images:
        full = str(p)
        rel = str(p.relative_to(root))
        issues = set(per_image.get(full, []) or per_image.get(rel, []) or per_image.get(p.name, []))
        if not issues:
            decisions[rel] = "keep"
            counts["keep"] += 1
            if move:
                _move(p, KEEP, root)
            continue
        if issues & HARD_ISSUES:
            decisions[rel] = "delete_candidate"
            counts["delete_candidate"] += 1
            if move:
                _move(p, DELETE, root)
        else:
            decisions[rel] = "review"
            counts["review"] += 1
            if move:
                _move(p, REVIEW, root)
    return {"counts": counts, "decisions": decisions}


def _move(p: Path, dest_root: Path, src_root: Path):
    """Move p into dest_root, preserving relative structure under source name."""
    rel = p.relative_to(src_root)
    dest = dest_root / p.parent.name / p.name
    dest.parent.mkdir(parents=True, exist_ok=True)
    shutil.move(str(p), str(dest))




def _slice_dir(base_dir: str, max_files: int) -> str:
    """Copy first N images into a temp dir for CleanVision to audit."""
    import shutil
    import tempfile
    from pathlib import Path

    src = Path(base_dir)
    tmp = Path(tempfile.mkdtemp(prefix="cvslice_"))
    exts = {".jpg", ".jpeg", ".png", ".webp", ".gif", ".heic", ".tif", ".bmp"}
    n = 0
    for p in src.rglob("*"):
        if p.is_file() and p.suffix.lower() in exts:
            shutil.copy2(p, tmp / p.name)
            n += 1
            if n >= max_files:
                break
    print(f"_slice_dir: copied {n} files to {tmp}")
    return str(tmp)


@flow(name="photo-quality-scan")
def quality_scan(base_dir: str, move: bool = False, notify: bool = True, max_files: int = None):
    """Audit image quality with CleanVision; classify into keep/review/delete.

    max_files: if set, only audit the first N image files (slices huge folders
    into reviewable chunks — prevents OOM on 50K-file trees).
    """
    if max_files:  # 0/None = unlimited
        base_dir = _slice_dir(base_dir, max_files)
    audit = audit_folder(base_dir)
    print(f"Issue summary: {audit['summary']}")
    result = classify_and_sort(base_dir, audit["per_image"], move=move)
    print(f"Verdicts: {result['counts']}")
    return {"audit": audit["summary"], **result}
    if notify:
        notify_result(result["counts"], audit["summary"], base_dir)


if __name__ == "__main__":
    import sys
    d = sys.argv[1] if len(sys.argv) > 1 else "/tmp/cvtest"
    move = "--move" in sys.argv
    quality_scan(d, move=move)


@task
def notify_result(counts: dict, summary: list, base_dir: str):
    """Send batch summary via Apprise."""
    import apprise_helper

    hard = counts.get("delete_candidate", 0)
    soft = counts.get("review", 0)
    keep = counts.get("keep", 0)
    flagged = [s for s in summary if s["num_images"] > 0]
    lines = "; ".join(f"{s[issue_type]}: {s[num_images]}" for s in flagged) or "none"
    body = (
        f"Scanned: {base_dir}\n"
        f"Keep: {keep} | Review: {soft} | Delete-candidates: {hard}\n"
        f"Issues: {lines}\n"
        f"Review: http://192.168.20.13:8092/review"
    )
    apprise_helper.notify("📸 photo-pipeline batch complete", body)
    return True