evaluation.py

from contextlib import contextmanager
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import Set
import os
import torch
from cleanfid import fid as CleanFID
from prettytable import PrettyTable
from pytorch_fid import fid_score as PytorchFID
from PIL import Image
from torch.utils.data import Dataset
from torchmetrics.image import PeakSignalNoiseRatio
from torchmetrics.image import StructuralSimilarityIndexMeasure
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
from torchvision import transforms
from tqdm import tqdm

IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png'}


class EvalDataset(Dataset):
    def __init__(self, gt_folder, pred_folder, height=1024):
        self.gt_folder = gt_folder
        self.pred_folder = pred_folder
        self.height = height
        self.data = self.prepare_data()
        self.to_tensor = transforms.ToTensor()

    def extract_id_from_filename(self, filename):
        if "inshop" in filename:
            filename = filename.split(".")[0]
            return filename

        # find first number in filename
        start_i = None
        for i, c in enumerate(filename):
            if c.isdigit():
                start_i = i
                break
        if start_i is None:
            assert False, f"Cannot find number in filename {filename}"
        return filename[start_i:start_i+8]

    def prepare_data(self):
        gt_files = scan_files_in_dir(self.gt_folder, postfix=IMAGE_EXTENSIONS)
        gt_dict = {self.extract_id_from_filename(
            file.name): file for file in gt_files}
        pred_files = scan_files_in_dir(
            self.pred_folder, postfix=IMAGE_EXTENSIONS)

        tuples = []
        for pred_file in pred_files:
            pred_id = self.extract_id_from_filename(pred_file.name)
            if pred_id not in gt_dict:
                print(f"Cannot find gt file for {pred_file}")
            else:
                tuples.append((gt_dict[pred_id].path, pred_file.path))
        return tuples

    def resize(self, img):
        w, h = img.size
        new_w = int(w * self.height / h)
        return img.resize((new_w, self.height), Image.LANCZOS)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        gt_path, pred_path = self.data[idx]
        gt, pred = self.resize(Image.open(gt_path)), self.resize(
            Image.open(pred_path))
        if gt.height != self.height:
            gt = self.resize(gt)
        if pred.height != self.height:
            pred = self.resize(pred)
        gt = self.to_tensor(gt)
        pred = self.to_tensor(pred)
        return gt, pred

def scan_files_in_dir(directory, postfix: Set[str] = None, progress_bar: tqdm = None) -> list:
    file_list = []
    progress_bar = tqdm(total=0, desc=f"Scanning",
                        ncols=100) if progress_bar is None else progress_bar
    for entry in os.scandir(directory):
        if entry.is_file():
            if postfix is None or os.path.splitext(entry.path)[1].lower() in postfix:
                file_list.append(entry)
                progress_bar.total += 1
                progress_bar.update(1)
        elif entry.is_dir():
            file_list += scan_files_in_dir(entry.path,
                                           postfix=postfix, progress_bar=progress_bar)
    return file_list


def first_image_path(directory):
    """Return a deterministic image sample from a possibly nested folder."""
    files = scan_files_in_dir(directory, postfix=IMAGE_EXTENSIONS)
    if not files:
        raise RuntimeError(f"No images found in {directory}")
    return min(file.path for file in files)


@contextmanager
def flattened_image_dir(directory):
    """Expose recursive images in a temporary flat directory for PyTorch-FID."""
    files = sorted(
        scan_files_in_dir(directory, postfix=IMAGE_EXTENSIONS),
        key=lambda file: file.path,
    )
    if not files:
        raise RuntimeError(f"No images found in {directory}")

    with TemporaryDirectory(prefix="decovton-eval-") as temp_dir:
        temp_path = Path(temp_dir)
        for index, file in enumerate(files):
            suffix = Path(file.name).suffix.lower()
            (temp_path / f"{index:08d}{suffix}").symlink_to(
                Path(file.path).resolve()
            )
        yield temp_dir


def copy_resize_gt(gt_folder, height, width):
    new_folder = f"{gt_folder}_{height}"
    if not os.path.exists(new_folder):
        os.makedirs(new_folder, exist_ok=True)
    gt_files = scan_files_in_dir(gt_folder, postfix=IMAGE_EXTENSIONS)
    for file in tqdm(gt_files, desc="Resizing GT images"):
        relative_path = os.path.relpath(file.path, gt_folder)
        output_path = os.path.join(new_folder, relative_path)
        if os.path.exists(output_path):
            continue
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        with Image.open(file.path) as source:
            img = source.resize((width, height), Image.LANCZOS)
            img.save(output_path)
    return new_folder

@torch.no_grad()
def psnr(dataloader):
    psnr_score = 0
    psnr = PeakSignalNoiseRatio(data_range=1.0).to("cuda")
    for gt, pred in tqdm(dataloader, desc="Calculating PSNR"):
        batch_size = gt.size(0)
        gt, pred = gt.to("cuda"), pred.to("cuda")
        psnr_score += psnr(pred, gt) * batch_size
    return psnr_score / len(dataloader.dataset)

@torch.no_grad()
def ssim(dataloader):
    ssim_score = 0
    ssim = StructuralSimilarityIndexMeasure(data_range=1.0).to("cuda")
    for gt, pred in tqdm(dataloader, desc="Calculating SSIM"):
        batch_size = gt.size(0)
        gt, pred = gt.to("cuda"), pred.to("cuda")
        ssim_score += ssim(pred, gt) * batch_size
    return ssim_score / len(dataloader.dataset)

@torch.no_grad()
def lpips(dataloader):
    lpips_score = LearnedPerceptualImagePatchSimilarity(
        net_type='squeeze').to("cuda")
    score = 0
    for gt, pred in tqdm(dataloader, desc="Calculating LPIPS"):
        batch_size = gt.size(0)
        pred = pred.to("cuda")
        gt = gt.to("cuda")
        # LPIPS needs the images to be in the [-1, 1] range.
        gt = (gt * 2) - 1
        pred = (pred * 2) - 1
        score += lpips_score(gt, pred) * batch_size
    return score / len(dataloader.dataset)

def eval(args):
    # Check gt_folder has images with target height, resize if not
    with Image.open(first_image_path(args.pred_folder)) as pred_sample:
        pred_size = pred_sample.size
    with Image.open(first_image_path(args.gt_folder)) as gt_sample:
        gt_size = gt_sample.size
    if pred_size[1] != gt_size[1]:
        title = "--"*30 + \
            f"Resizing GT Images to height {pred_size[1]}" + "--"*30
        print(title)
        args.gt_folder = copy_resize_gt(
            args.gt_folder, pred_size[1], pred_size[0]
        )
        print("-"*len(title))

    # Calculate Metrics
    header = []
    row = []
    header += ["PyTorch-FID"]
    with flattened_image_dir(args.gt_folder) as fid_gt_folder, \
            flattened_image_dir(args.pred_folder) as fid_pred_folder:
        pytorch_fid_ = PytorchFID.calculate_fid_given_paths(
            [fid_gt_folder, fid_pred_folder], batch_size=args.batch_size,
            device="cuda", dims=2048, num_workers=args.num_workers)
        row += ["{:.4f}".format(pytorch_fid_)]
        print("PyTorch-FID: {:.4f}".format(pytorch_fid_))
        header += ["Clean-FID", "Clean-KID"]
        clean_fid_ = CleanFID.compute_fid(fid_gt_folder, fid_pred_folder)
        print("Clean-FID: {:.4f}".format(clean_fid_))
        row += ["{:.4f}".format(clean_fid_)]
        clean_kid_ = CleanFID.compute_kid(
            fid_gt_folder, fid_pred_folder
        ) * 1000
        print("Clean-KID: {:.4f}".format(clean_kid_))
        row += ["{:.4f}".format(clean_kid_)]
    if args.paired:
        dataset = EvalDataset(args.gt_folder, args.pred_folder, pred_size[1])
        if not dataset.data:
            raise RuntimeError(
                "No matching ground-truth/prediction image pairs were found."
            )
        dataloader = torch.utils.data.DataLoader(
            dataset, batch_size=args.batch_size, num_workers=args.num_workers,
            shuffle=False, drop_last=False
        )
        header += ["PSNR", "SSIM", "LPIPS"]
        psnr_ = psnr(dataloader).item()
        print("PSNR: {:.4f}".format(psnr_))
        ssim_ = ssim(dataloader).item()
        print("SSIM: {:.4f}".format(ssim_))
        lpips_ = lpips(dataloader).item()
        print("LPIPS: {:.4f}".format(lpips_))
        row += ["{:.4f}".format(psnr_), "{:.4f}".format(ssim_),
                "{:.4f}".format(lpips_)]
    # Print Results
    print("GT Folder  : ", args.gt_folder)
    print("Pred Folder: ", args.pred_folder)
    table = PrettyTable()
    table.field_names = header
    table.add_row(row)
    print(table)

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--gt_folder", type=str, required=True)
    parser.add_argument("--pred_folder", type=str, required=True)
    parser.add_argument("--paired", action="store_true")
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--num_workers", type=int, default=4)
    args = parser.parse_args()

    if args.gt_folder.endswith("/"):
        args.gt_folder = args.gt_folder[:-1]
    if args.pred_folder.endswith("/"):
        args.pred_folder = args.pred_folder[:-1]

    eval(args)