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