diff --git a/src/shortlist_top.py b/src/shortlist_top.py new file mode 100644 index 0000000..3473c3e --- /dev/null +++ b/src/shortlist_top.py @@ -0,0 +1,190 @@ +"""Shortlist the top N% of scored photos. + +Reads output.json, sorts by laion_score, picks the top X%, +and copies them to a destination folder. + +Usage: + python -m src.shortlist_top +""" + +from __future__ import annotations + +import json +import math +import os +import shutil +import sys + +# Add project root to path +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +# --------------------------------------------------------------------------- +# IO +# --------------------------------------------------------------------------- + + +def load_results(output_path: str) -> list[dict]: + """Load results from JSON file.""" + if not os.path.exists(output_path): + print(f" ERROR: {output_path} not found.") + sys.exit(1) + try: + with open(output_path, encoding="utf-8") as f: + data = json.load(f) + if not isinstance(data, list): + print(" ERROR: Expected a JSON array.") + sys.exit(1) + return data + except json.JSONDecodeError as e: + print(f" ERROR: Invalid JSON: {e}") + sys.exit(1) + + +# --------------------------------------------------------------------------- +# Display +# --------------------------------------------------------------------------- + + +def show_summary(results: list[dict]) -> None: + """Print a quick summary of scored photos.""" + scored = [r for r in results if "laion_score" in r] + print() + print("=" * 50) + print(" Scored Photos Summary") + print("=" * 50) + print(f"\n Total entries: {len(results)}") + print(f" With scores: {len(scored)}") + if scored: + scores = [r["laion_score"] for r in scored] + print(f" Score range: {min(scores):.2f} — {max(scores):.2f}") + print(f" Average: {sum(scores) / len(scores):.2f}") + print() + + +# --------------------------------------------------------------------------- +# Filtering & Copying +# --------------------------------------------------------------------------- + + +def pick_top_pct(results: list[dict], pct: float) -> list[dict]: + """Return the top N% of scored results by laion_score.""" + scored = [r for r in results if "laion_score" in r] + if not scored: + return [] + scored.sort(key=lambda r: r["laion_score"], reverse=True) + n = max(1, math.ceil(len(scored) * pct / 100)) + return scored[:n] + + +def copy_photos( + photos: list[dict], dest_dir: str, dry_run: bool = False +) -> tuple[int, int]: + """Copy photos to dest_dir. Returns (copied, skipped).""" + try: + os.makedirs(dest_dir, exist_ok=True) + except OSError as e: + print(f" ERROR: Could not create {dest_dir}: {e}") + return 0, 0 + + copied = 0 + skipped = 0 + + for entry in photos: + filepath = entry["filepath"] + filename = entry["filename"] + + if not os.path.exists(filepath): + print(f" Skip: {filepath} not found") + skipped += 1 + continue + + dest = os.path.join(dest_dir, filename) + if os.path.exists(dest): + base, ext = os.path.splitext(filename) + counter = 1 + while os.path.exists(dest): + dest = os.path.join(dest_dir, f"{base}_{counter}{ext}") + counter += 1 + + if not dry_run: + shutil.copy2(filepath, dest) + copied += 1 + + return copied, skipped + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + print() + print("=" * 50) + print(" Photo Judgers — Top N% Shortlist") + print("=" * 50) + + # Load results + output_path = input("\n Output JSON file [output.json]: ").strip() or "output.json" + results = load_results(output_path) + show_summary(results) + + # Ask for percentile + while True: + pct_str = input(" Top N% [10]: ").strip() + try: + pct = float(pct_str) if pct_str else 10.0 + if 0 < pct <= 100: + break + except ValueError: + pass + print(" Enter a number between 0 and 100.") + + # Calculate how many photos + scored = [r for r in results if "laion_score" in r] + n = max(1, math.ceil(len(scored) * pct / 100)) + print(f"\n Top {pct}% = {n} of {len(scored)} scored photos") + + # Show the cutoff score + scored_sorted = sorted(scored, key=lambda r: r["laion_score"], reverse=True) + cutoff = scored_sorted[n - 1]["laion_score"] + print(f" Cutoff score: {cutoff:.2f}") + + # Show sample filenames + print("\n Shortlisted photos:") + for entry in scored_sorted[:n]: + print(f" {entry['laion_score']:.2f} {entry['filename']}") + + # Destination folder + dest = input("\n Shortlist folder: ").strip() or "shortlist_top" + + # Dry-run option + dry = input(" Dry run? (don't copy) [y/N]: ").strip().lower() in ("y", "yes") + if dry: + print(" Dry run — no files will be copied.") + + # Confirm + print() + confirm = input(" Continue? [y/N]: ").strip().lower() + if confirm not in ("y", "yes"): + print(" Aborted.") + return + + # Pick and copy + chosen = pick_top_pct(results, pct) + print(f"\n Copying to {dest}...") + copied, skipped = copy_photos(chosen, dest, dry_run=dry) + + print() + print(f" Copied: {copied}") + if skipped: + print(f" Skipped: {skipped}") + if not dry: + print(f" Done. Shortlist: {os.path.abspath(dest)}") + else: + print(" (dry run — no files copied)") + + +if __name__ == "__main__": + main()