from sdeval.corrupt import AICorruptMetrics
import json
import os
from pathlib import Path
from PIL import Image
from tqdm import tqdm
import argparse
# build metrics
metrics = AICorruptMetrics()

# Define the directory containing .webp files and the output JSON file path
def infer_path_to_result(input_dir,output_json):
    DIRECTORY_PATH = Path(input_dir)
    OUTPUT_JSON_PATH = Path(output_json)
    # Function to calculate aesthetic score for a single image
    def calculate_score(image_path):
        try:
            score = metrics.score(image_path)
            return score
        except Exception as e:
            print(f"Error processing {image_path}: {e}")
            return None

    # Process all .webp files in the directory
    scores = {}
    all_images = list(DIRECTORY_PATH.rglob("*.webp"))
    for image_path in tqdm(all_images):
        image_path = str(image_path)
        score = calculate_score(image_path)
        if score is not None:
            scores[image_path] = score

    # Save scores to JSON
    with OUTPUT_JSON_PATH.open("w") as f:
        json.dump(scores, f, indent=4)

    print(f"Corruption scores saved to {OUTPUT_JSON_PATH}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Process aesthetic scores for images in a directory')
    parser.add_argument('--input_dir', type=str, help='Directory containing .webp files')
    parser.add_argument('--output_json', type=str, help='Output JSON file path')
    args = parser.parse_args()
    infer_path_to_result(args.input_dir, args.output_json)
