feat: add shortlist_top — pick top N% of scored photos and copy to folder
This commit is contained in:
@@ -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()
|
||||
Reference in New Issue
Block a user