eval.py

import os
import yaml
import torch
from pytorch_fid import fid_score
from src import set_seed

if __name__ == "__main__":
    with open("config/eval.yaml", "r") as f:
        config = yaml.safe_load(f)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")
    
    # OS
    os_config = config["os"]
    set_seed(os_config["seed"])
    
    # Data
    data_config = config["data"]
    path_real_images = data_config["path_real_images"]
    path_generated_images = data_config["path_generated_images"]

    if not os.path.isdir(path_real_images):
        raise NotADirectoryError(
            f"Real image path does not exist or is not a directory: {path_real_images}")
    if not os.path.isdir(path_generated_images):
        raise NotADirectoryError(
            f"Generated image path does not exist or is not a directory: {path_generated_images}")

    # Inference
    inference_config = config["inference"]
    batch_size = inference_config["batch_size"]
    dims = inference_config["dims"]
    
    print("Calculating FID score...")
    try:
        fid_value = fid_score.calculate_fid_given_paths(
            paths=[path_real_images, path_generated_images],
            batch_size=batch_size,
            device=device,
            dims=dims,
            num_workers=8
        )
        print(f"Calculated FID score: {fid_value}")

    except Exception as e:
        print(f"An error occurred while calculating FID: {e}")