chore: remove root scorer.py (moved to src/)
This commit is contained in:
@@ -1,283 +0,0 @@
|
|||||||
"""
|
|
||||||
Photo Judgers — Scorer
|
|
||||||
|
|
||||||
Interactive scorer with menu for selecting which model to run.
|
|
||||||
Tracks progress by reading existing JSON output to avoid re-scoring.
|
|
||||||
|
|
||||||
Usage:
|
|
||||||
python scorer.py
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import glob
|
|
||||||
import json
|
|
||||||
import os
|
|
||||||
import sys
|
|
||||||
from datetime import datetime, timezone
|
|
||||||
from typing import TYPE_CHECKING
|
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
|
||||||
from types import NoneType
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# Scorer Base
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
class BaseScorer:
|
|
||||||
"""Base class for all scorers."""
|
|
||||||
|
|
||||||
name: str = "Base"
|
|
||||||
description: str = ""
|
|
||||||
|
|
||||||
def load(self) -> None:
|
|
||||||
"""Load model weights / dependencies. Override in subclasses."""
|
|
||||||
|
|
||||||
def score(self, image_path: str) -> dict | None:
|
|
||||||
"""
|
|
||||||
Score a single image.
|
|
||||||
|
|
||||||
Returns a dict with scorer-specific fields, or None on failure.
|
|
||||||
"""
|
|
||||||
raise NotImplementedError
|
|
||||||
|
|
||||||
def unload(self) -> None:
|
|
||||||
"""Clean up resources. Override in subclasses."""
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# LAION Scorer
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
class LaionScorer(BaseScorer):
|
|
||||||
name = "LAION Aesthetic Predictor V2"
|
|
||||||
description = "CLIP-based aesthetic scoring (0-10 scale)"
|
|
||||||
|
|
||||||
def __init__(self) -> None:
|
|
||||||
self.model: object = None # type: ignore[assignment]
|
|
||||||
self.processor: object = None # type: ignore[assignment]
|
|
||||||
self._version = "v2"
|
|
||||||
|
|
||||||
def load(self) -> None:
|
|
||||||
try:
|
|
||||||
import torch # noqa: F401
|
|
||||||
from simple_aesthetics_predictor import AestheticsPredictorV1 # noqa: F401
|
|
||||||
from transformers import CLIPProcessor # noqa: F401
|
|
||||||
|
|
||||||
print(" Loading LAION V2 model... (first run downloads ~2GB)")
|
|
||||||
from simple_aesthetics_predictor import AestheticsPredictorV1
|
|
||||||
from transformers import CLIPProcessor
|
|
||||||
|
|
||||||
self.model = AestheticsPredictorV1.from_pretrained(
|
|
||||||
"shunk031/aesthetics-predictor-v2-vit-large-patch14"
|
|
||||||
)
|
|
||||||
self.processor = CLIPProcessor.from_pretrained(
|
|
||||||
"shunk031/aesthetics-predictor-v2-vit-large-patch14"
|
|
||||||
)
|
|
||||||
self.model.eval()
|
|
||||||
print(" Model loaded.")
|
|
||||||
except ImportError:
|
|
||||||
print(" ERROR: Required packages not installed.")
|
|
||||||
print(
|
|
||||||
" Run: pip install simple-aesthetics-predictor transformers "
|
|
||||||
"torch torchvision Pillow tqdm"
|
|
||||||
)
|
|
||||||
sys.exit(1)
|
|
||||||
|
|
||||||
def score(self, image_path: str) -> dict | None:
|
|
||||||
try:
|
|
||||||
import torch # noqa: F401
|
|
||||||
from PIL import Image
|
|
||||||
|
|
||||||
assert self.processor is not None, "Model not loaded. Call load() first."
|
|
||||||
assert self.model is not None, "Model not loaded. Call load() first."
|
|
||||||
|
|
||||||
image = Image.open(image_path).convert("RGB")
|
|
||||||
inputs = self.processor(images=image, return_tensors="pt")
|
|
||||||
|
|
||||||
with torch.no_grad():
|
|
||||||
outputs = self.model(**inputs)
|
|
||||||
score = outputs.logits.squeeze().item()
|
|
||||||
|
|
||||||
return {
|
|
||||||
"laion_score": round(score, 2),
|
|
||||||
"laion_version": self._version,
|
|
||||||
}
|
|
||||||
except Exception as e:
|
|
||||||
print(f" Warning: Failed to score {image_path}: {e}")
|
|
||||||
return None
|
|
||||||
|
|
||||||
def unload(self) -> None:
|
|
||||||
self.model = None # type: ignore[assignment]
|
|
||||||
self.processor = None # type: ignore[assignment]
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# Registry
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
SCORERS: list[type[BaseScorer]] = [
|
|
||||||
LaionScorer,
|
|
||||||
# Add new scorers here:
|
|
||||||
# NimaScorer,
|
|
||||||
# MusiqScorer,
|
|
||||||
# BrisqueScorer,
|
|
||||||
]
|
|
||||||
|
|
||||||
|
|
||||||
def show_scorer_menu() -> type[BaseScorer]:
|
|
||||||
"""Display scorer selection menu and return the chosen scorer class."""
|
|
||||||
print()
|
|
||||||
print("=" * 50)
|
|
||||||
print(" Photo Judgers — Select Scorer")
|
|
||||||
print("=" * 50)
|
|
||||||
print()
|
|
||||||
for i, scorer_cls in enumerate(SCORERS, 1):
|
|
||||||
print(f" {i}. {scorer_cls.name}")
|
|
||||||
print(f" {scorer_cls.description}")
|
|
||||||
print()
|
|
||||||
|
|
||||||
while True:
|
|
||||||
choice = input(f" Choose [1-{len(SCORERS)}]: ").strip()
|
|
||||||
try:
|
|
||||||
idx = int(choice) - 1
|
|
||||||
if 0 <= idx < len(SCORERS):
|
|
||||||
return SCORERS[idx]
|
|
||||||
except ValueError:
|
|
||||||
pass
|
|
||||||
print(" Invalid choice. Try again.")
|
|
||||||
|
|
||||||
|
|
||||||
def ask_folder() -> str:
|
|
||||||
"""Prompt user for the input folder path."""
|
|
||||||
while True:
|
|
||||||
folder = input("\n Enter folder path to score: ").strip().strip('"\'')
|
|
||||||
if not folder:
|
|
||||||
print(" Path cannot be empty.")
|
|
||||||
continue
|
|
||||||
if os.path.isdir(folder):
|
|
||||||
return folder
|
|
||||||
print(f" Folder not found: {folder}")
|
|
||||||
|
|
||||||
|
|
||||||
def ask_output() -> str:
|
|
||||||
"""Prompt user for the output JSON file path."""
|
|
||||||
default = "output.json"
|
|
||||||
answer = input(f"\n Output JSON file [{default}]: ").strip()
|
|
||||||
return answer if answer else default
|
|
||||||
|
|
||||||
|
|
||||||
def load_existing_results(output_path: str) -> dict:
|
|
||||||
"""Load existing results from JSON file. Returns {filepath: result}."""
|
|
||||||
if os.path.exists(output_path):
|
|
||||||
try:
|
|
||||||
with open(output_path, encoding="utf-8") as f:
|
|
||||||
data = json.load(f)
|
|
||||||
if isinstance(data, list):
|
|
||||||
return {entry["filepath"]: entry for entry in data}
|
|
||||||
except (json.JSONDecodeError, OSError) as e:
|
|
||||||
print(f" Warning: Could not read existing results: {e}")
|
|
||||||
return {}
|
|
||||||
|
|
||||||
|
|
||||||
def get_image_files(folder: str) -> list[str]:
|
|
||||||
"""Get all image files from folder (recursively)."""
|
|
||||||
extensions = {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".webp", ".gif"}
|
|
||||||
files: list[str] = []
|
|
||||||
for ext in extensions:
|
|
||||||
files.extend(glob.glob(os.path.join(folder, f"**/*{ext}"), recursive=True))
|
|
||||||
files.extend(glob.glob(os.path.join(folder, f"*{ext}")))
|
|
||||||
return sorted(set(files))
|
|
||||||
|
|
||||||
|
|
||||||
def run_scorer(scorer_cls: type[BaseScorer], folder: str, output_path: str) -> None:
|
|
||||||
"""Run the selected scorer on all images in the folder."""
|
|
||||||
print(f"\n Scorer: {scorer_cls.name}")
|
|
||||||
print(f" Folder: {folder}")
|
|
||||||
print(f" Output: {output_path}")
|
|
||||||
print()
|
|
||||||
|
|
||||||
# Load existing results
|
|
||||||
existing = load_existing_results(output_path)
|
|
||||||
print(f" Already scored: {len(existing)} images")
|
|
||||||
|
|
||||||
# Get all image files
|
|
||||||
all_files = get_image_files(folder)
|
|
||||||
remaining = [f for f in all_files if f not in existing]
|
|
||||||
print(f" Remaining to score: {len(remaining)} images")
|
|
||||||
print()
|
|
||||||
|
|
||||||
if not remaining:
|
|
||||||
print(" Nothing new to score. Done!")
|
|
||||||
return
|
|
||||||
|
|
||||||
# Load the model
|
|
||||||
scorer = scorer_cls()
|
|
||||||
scorer.load()
|
|
||||||
|
|
||||||
# Process images
|
|
||||||
results = list(existing.values()) # Start with existing
|
|
||||||
skipped = 0
|
|
||||||
|
|
||||||
from tqdm import tqdm
|
|
||||||
|
|
||||||
for image_path in tqdm(remaining, desc="Scoring"):
|
|
||||||
result = scorer.score(image_path)
|
|
||||||
if result:
|
|
||||||
entry = {
|
|
||||||
"filename": os.path.basename(image_path),
|
|
||||||
"filepath": image_path,
|
|
||||||
"subtype": scorer_cls.name.lower().replace(" ", "_"),
|
|
||||||
**result,
|
|
||||||
"scoring_date": datetime.now(timezone.utc).isoformat(),
|
|
||||||
}
|
|
||||||
results.append(entry)
|
|
||||||
else:
|
|
||||||
skipped += 1
|
|
||||||
|
|
||||||
scorer.unload()
|
|
||||||
|
|
||||||
# Save results
|
|
||||||
try:
|
|
||||||
with open(output_path, "w", encoding="utf-8") as f:
|
|
||||||
json.dump(results, f, indent=2)
|
|
||||||
except OSError as e:
|
|
||||||
print(f" ERROR: Could not write output file: {e}")
|
|
||||||
return
|
|
||||||
|
|
||||||
print()
|
|
||||||
print(f" Scored: {len(remaining) - skipped} images")
|
|
||||||
if skipped:
|
|
||||||
print(f" Skipped (failed): {skipped} images")
|
|
||||||
print(f" Total in output: {len(results)} images")
|
|
||||||
print(f" Saved to: {output_path}")
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# Main
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
def main() -> None:
|
|
||||||
print()
|
|
||||||
print("=" * 50)
|
|
||||||
print(" Photo Judgers")
|
|
||||||
print("=" * 50)
|
|
||||||
|
|
||||||
# Select scorer
|
|
||||||
scorer_cls = show_scorer_menu()
|
|
||||||
|
|
||||||
# Get folder
|
|
||||||
folder = ask_folder()
|
|
||||||
|
|
||||||
# Get output file
|
|
||||||
output_path = ask_output()
|
|
||||||
|
|
||||||
# Run
|
|
||||||
run_scorer(scorer_cls, folder, output_path)
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
Reference in New Issue
Block a user