104 lines
3.6 KiB
Python
104 lines
3.6 KiB
Python
"""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))
|
|
|
|
|
|
@flow(name="photo-quality-scan")
|
|
def quality_scan(base_dir: str, move: bool = False):
|
|
"""Audit image quality with CleanVision; classify into keep/review/delete."""
|
|
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 __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)
|