test_vae.py

import argparse
import os
from PIL import Image
import torch
from torchvision.transforms import Resize, ToTensor
from diffusers import AutoencoderKL
from pytorch_fid import fid_score
from skimage.metrics import peak_signal_noise_ratio as psnr
import lpips
from tqdm import tqdm
from torchvision import transforms

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def load_images(folder_path):
    images = []
    for filename in os.listdir(folder_path):
        if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
            img_path = os.path.join(folder_path, filename)
            images.append(img_path)
    return images


def paramiter_count(model):
    state_dict = model.state_dict()
    paramiter_count = 0
    for key in state_dict:
        paramiter_count += torch.numel(state_dict[key])
    return int(paramiter_count)


def calculate_metrics(vae, images, max_imgs=-1, save_output=False):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    vae = vae.to(device)
    lpips_model = lpips.LPIPS(net='alex').to(device)

    rfid_scores = []
    psnr_scores = []
    lpips_scores = []

    # transform = transforms.Compose([
    #     transforms.Resize(256, antialias=True),
    #     transforms.CenterCrop(256)
    # ])
    # needs values between -1 and 1
    to_tensor = ToTensor()
    
    # remove _reconstructed.png files
    images = [img for img in images if not img.endswith("_reconstructed.png")]

    if max_imgs > 0 and len(images) > max_imgs:
        images = images[:max_imgs]

    for img_path in tqdm(images):
        try:
            img = Image.open(img_path).convert('RGB')
            # img_tensor = to_tensor(transform(img)).unsqueeze(0).to(device)
            img_tensor = to_tensor(img).unsqueeze(0).to(device)
            img_tensor = 2 * img_tensor - 1
            # if width or height is not divisible by 8, crop it
            if img_tensor.shape[2] % 8 != 0 or img_tensor.shape[3] % 8 != 0:
                img_tensor = img_tensor[:, :, :img_tensor.shape[2] // 8 * 8, :img_tensor.shape[3] // 8 * 8]

        except Exception as e:
            print(f"Error processing {img_path}: {e}")
            continue


        with torch.no_grad():
            reconstructed = vae.decode(vae.encode(img_tensor).latent_dist.sample()).sample

        # Calculate rFID
        # rfid = fid_score.calculate_frechet_distance(vae, img_tensor, reconstructed)
        # rfid_scores.append(rfid)

        # Calculate PSNR
        psnr_val = psnr(img_tensor.cpu().numpy(), reconstructed.cpu().numpy())
        psnr_scores.append(psnr_val)

        # Calculate LPIPS
        lpips_val = lpips_model(img_tensor, reconstructed).item()
        lpips_scores.append(lpips_val)

    # avg_rfid = sum(rfid_scores) / len(rfid_scores)
    avg_rfid = 0
    avg_psnr = sum(psnr_scores) / len(psnr_scores)
    avg_lpips = sum(lpips_scores) / len(lpips_scores)
    
    if save_output:
        filename_no_ext = os.path.splitext(os.path.basename(img_path))[0]
        folder = os.path.dirname(img_path)
        save_path = os.path.join(folder, filename_no_ext + "_reconstructed.png")
        reconstructed = (reconstructed + 1) / 2
        reconstructed = reconstructed.clamp(0, 1)
        reconstructed = transforms.ToPILImage()(reconstructed[0].cpu())
        reconstructed.save(save_path)

    return avg_rfid, avg_psnr, avg_lpips


def main():
    parser = argparse.ArgumentParser(description="Calculate average rFID, PSNR, and LPIPS for VAE reconstructions")
    parser.add_argument("--vae_path", type=str, required=True, help="Path to the VAE model")
    parser.add_argument("--image_folder", type=str, required=True, help="Path to the folder containing images")
    parser.add_argument("--max_imgs", type=int, default=-1, help="Max num of images. Default is -1 for all images.")
    # boolean store true
    parser.add_argument("--save_output", action="store_true", help="Save the output images")
    args = parser.parse_args()

    if  os.path.isfile(args.vae_path):
        vae = AutoencoderKL.from_single_file(args.vae_path)
    else:
        try:
            vae = AutoencoderKL.from_pretrained(args.vae_path)
        except:
            vae = AutoencoderKL.from_pretrained(args.vae_path, subfolder="vae")
    vae.eval()
    vae = vae.to(device)
    print(f"Model has {paramiter_count(vae)} parameters")
    images = load_images(args.image_folder)

    avg_rfid, avg_psnr, avg_lpips = calculate_metrics(vae, images, args.max_imgs, args.save_output)

    # print(f"Average rFID: {avg_rfid}")
    print(f"Average PSNR: {avg_psnr}")
    print(f"Average LPIPS: {avg_lpips}")


if __name__ == "__main__":
    main()