results_benchmark.py

"""Code for computation of PLCC, SRCC and KRCC between
    PIQ metrics predictions and ground truth scores from MOS databases.
"""
import piq
import tqdm
import torch
import types
import argparse
import functools
import torchvision

import pandas as pd

from typing import List, Callable, Tuple, Optional
from pathlib import Path
from skimage.io import imread
from scipy.stats import spearmanr, kendalltau
from torch.utils.data import DataLoader, Dataset
from dataclasses import dataclass
from torch import nn
from itertools import chain
from scipy.io import loadmat


OPENAI_CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)
OPENAI_CLIP_STD = (0.26862954, 0.26130258, 0.27577711)


@dataclass
class Metric:
    name: str
    functor: Callable
    category: str  # FR - full-reference, NR - no-reference, DB - distribution-based

    def __post_init__(self):
        valid_categories = {'FR', 'NR', 'DB'}
        assert self.category in valid_categories, f'Provided category [{self.category}] is invalid. ' \
                                                  f'Provide one of: {valid_categories}'


torch.multiprocessing.set_sharing_strategy("file_system")

METRICS = {
    # Full-reference
    "PSNR": Metric(name="PSNR", functor=functools.partial(piq.psnr, data_range=255., reduction='none'), category='FR'),
    "SSIM": Metric(name="SSIM", functor=functools.partial(piq.ssim, data_range=255., reduction='none'), category='FR'),
    "MS-SSIM": Metric(name="MS-SSIM", functor=functools.partial(piq.multi_scale_ssim, data_range=255.,
                                                                reduction='none'), category='FR'),
    "IW-SSIM": Metric(name="IW-SSIM", functor=functools.partial(piq.information_weighted_ssim, data_range=255.,
                                                                reduction='none'), category='FR'),
    "VIFp": Metric(name="VIFp", functor=functools.partial(piq.vif_p, data_range=255., reduction='none'), category='FR'),
    "GMSD": Metric(name="GMSD", functor=functools.partial(piq.gmsd, data_range=255., reduction='none'), category='FR'),
    "MS-GMSD": Metric(name="MS-GMSD", functor=functools.partial(piq.multi_scale_gmsd, data_range=255.,
                                                                reduction='none'), category='FR'),
    "MS-GMSDc": Metric(name="MS-GMSDc", functor=functools.partial(piq.multi_scale_gmsd, data_range=255.,
                                                                  chromatic=True, reduction='none'), category='FR'),
    "FSIM": Metric(name="FSIM", functor=functools.partial(piq.fsim, data_range=255.,
                                                          chromatic=False, reduction='none'), category='FR'),
    "FSIMc": Metric(name="FSIMc", functor=functools.partial(piq.fsim, data_range=255.,
                                                            chromatic=True, reduction='none'), category='FR'),
    "VSI": Metric(name="VSI", functor=functools.partial(piq.vsi, data_range=255., reduction='none'), category='FR'),
    "SR-SIM": Metric(name="SR-SIM", functor=functools.partial(piq.srsim, data_range=255.,
                                                              reduction='none'), category='FR'),
    "SR-SIMc": Metric(name="SR-SIM", functor=functools.partial(piq.srsim, data_range=255.,
                                                               reduction='none', chromatic=True), category='FR'),
    "HaarPSI": Metric(name="HaarPSI", functor=functools.partial(piq.haarpsi, data_range=255.,
                                                                reduction='none'), category='FR'),
    "MDSI": Metric(name="MDSI", functor=functools.partial(piq.mdsi, data_range=255., reduction='none'), category='FR'),
    "LPIPS-vgg": Metric(name="LPIPS-vgg", functor=piq.LPIPS(replace_pooling=False, reduction='none'), category='FR'),
    "DISTS": Metric(name="DISTS", functor=piq.DISTS(reduction='none'), category='FR'),
    "PieAPP": Metric(name="PieAPP", functor=piq.PieAPP(data_range=255., reduction='none'), category='FR'),
    "Content": Metric(name="Content", functor=piq.ContentLoss(reduction='none'), category='FR'),
    "Style": Metric(name="Style", functor=piq.StyleLoss(reduction='none'), category='FR'),
    "DSS": Metric(name="DSS", functor=functools.partial(piq.dss, data_range=255., reduction='none'), category='FR'),

    # No-reference
    "BRISQUE": Metric(
        name="BRISQUE",
        functor=functools.partial(piq.brisque, data_range=255.0, reduction="none"),
        category="NR",
    ),
    "CLIPIQA": Metric(name="CLIPIQA", functor=piq.CLIPIQA(data_range=255), category="NR"),

    # Distribution-based
    "IS": Metric(name="IS", functor=piq.IS(distance="l1"), category="DB"),
    "FID": Metric(name="FID", functor=piq.FID(), category="DB"),
    "GS": Metric(name="GS", functor=piq.GS(), category="DB"),
    "KID": Metric(name="KID", functor=piq.KID(), category="DB"),
    "MSID": Metric(name="MSID", functor=piq.MSID(), category="DB"),
    "PR": Metric(name="PR", functor=piq.PR(), category="DB"),
}

METRIC_CATEGORIES = {
    cat: [k for k, v in METRICS.items() if v.category == cat]
    for cat in ["FR", "NR", "DB"]
}


class TID2013(Dataset):
    r"""A class to evaluate on the KADID10k dataset.
    Note that the class is callable. The values are returned as a result of calling the __getitem__ method.

    Args:
        root: Root directory path.
    Returns:
        x: image with some kind of distortion in [0, 1] range
        y: image without distortion in [0, 1] range
        score: MOS score for this pair of images
    """
    _filename = "mos_with_names.txt"

    def __init__(self, root: Path = "datasets/tid2013") -> None:
        assert (
            root.exists()
        ), "You need to download TID2013 dataset first. Check http://www.ponomarenko.info/tid2013"

        df = pd.read_csv(
            root / self._filename, sep=" ", names=["score", "dist_img"], header=None
        )
        df["ref_img"] = df["dist_img"].apply(
            lambda x: f"reference_images/{(x[:3] + x[-4:]).upper()}"
        )
        df["dist_img"] = df["dist_img"].apply(lambda x: f"distorted_images/{x}")

        self.scores = df["score"].to_numpy()
        self.df = df[["dist_img", "ref_img", "score"]]
        self.root = root

    def __getitem__(self, index) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        x_path = self.root / self.df.iloc[index][0]
        y_path = self.root / self.df.iloc[index][1]
        score = self.scores[index]

        # Load image and ref, convert to tensor.
        x = torch.tensor(imread(x_path), dtype=torch.float32).permute(2, 0, 1)
        y = torch.tensor(imread(y_path), dtype=torch.float32).permute(2, 0, 1)

        return x, y, score

    def __len__(self) -> int:
        return len(self.df)


class KADID10k(TID2013):
    r"""A class to evaluate on the KADID10k dataset.
    One can get the dataset via the direct link: https://datasets.vqa.mmsp-kn.de/archives/kadid10k.zip.
    Note that the class is callable. The values are returned as a result of calling the __getitem__ method.

    Args:
        root: Root directory path.
    Returns:
        x: image with some kind of distortion in [0, 1] range
        y: image without distortion in [0, 1] range
        score: MOS score for this pair of images
    """
    _filename = "dmos.csv"

    def __init__(self, root: Path = "datasets/kadid10k") -> None:
        assert root.exists(), \
            "You need to download KADID10K dataset first. " \
            "Check http://database.mmsp-kn.de/kadid-10k-database.html " \
            "or download via the direct link https://datasets.vqa.mmsp-kn.de/archives/kadid10k.zip"

        # Read file mith DMOS
        self.df = pd.read_csv(root / self._filename)
        self.df.rename(columns={"dmos": "score"}, inplace=True)
        self.scores = self.df["score"].to_numpy()
        self.df = self.df[["dist_img", "ref_img", "score"]]

        self.root = root / "images"


class PIPAL(TID2013):
    r"""A class to evaluate on the train set of the PIPAL dataset.
    Note that the class is callable. The values are returned as a result of calling the __getitem__ method.

    Args:
        root: Root directory path.
    Returns:
        x: image with some kind of distortion in [0, 1] range
        y: image without distortion in [0, 1] range
        score: MOS score for this pair of images
    """
    def __init__(self, root: Path = Path("data/raw/pipal")) -> None:
        assert root.exists(), \
            "You need to download PIPAL dataset. Check https://www.jasongt.com/projectpages/pipal.html"

        assert (root / "Train_Dist").exists(), \
            "Please place all distorted files into single folder named `Train_Dist`."

        # Read files with labels and merge them into single DF
        dfs = []
        for filename in (root / "Train_Label").glob("*.txt"):
            df = pd.read_csv(filename, index_col=None, header=None, names=['dist_img', 'score'])
            dfs.append(df)

        df = pd.concat(dfs, axis=0, ignore_index=True)

        df["ref_img"] = df["dist_img"].apply(lambda x: f"Train_Ref/{x[:5] + x[-4:]}")
        df["dist_img"] = df["dist_img"].apply(lambda x: f"Train_Dist/{x}")

        self.scores = df["score"].to_numpy()
        self.df = df[["dist_img", "ref_img", "score"]]
        self.root = root


class KonIQ10k(Dataset):
    r"""A class to evaluate on the train/test/train+test the KonIQ dataset.
    http://database.mmsp-kn.de/koniq-10k-database.html

    Note that the class is callable. The values are returned as a result of calling the __getitem__ method.

    Args:
        root: Root directory path.
        transform: Something that pre-processes dataset images.
        subset: Which part of the dataset to use.
            Options: "train", "test", "all". "all" means train + test.
    Returns:
        x: image with some kind of distortion in [0, 1] range.
        y: dummy variable, used for compatibility with FR datasets.
        score: MOS score of the image quality.
    """
    _filename_scores = "koniq10k_scores.csv"
    _filename_dists = "koniq10k_distributions_sets.csv"

    def __init__(self, root: Path, transform: Optional[Callable] = None, subset: str = 'test') -> None:
        supported_subsets = ["train", "test", "all"]
        assert subset in supported_subsets, f"Unknown subset [{subset}], choose one of {supported_subsets}."
        assert root.exists(), "You need to download KonIQ-10k dataset first."
        
        self.root = root
        self.initial_image_size = "1024x768"
        df1 = pd.read_csv(root / self._filename_scores)
        df2 = pd.read_csv(root / self._filename_dists)
        self.df = df1.merge(df2, on=["image_name"])
        if not subset == "all":
            self.df = self.df[self.df.set == subset].reset_index()

        self.df["image_name"] = self.df["image_name"].apply(
            lambda x: f"{self.initial_image_size}/{x}"
        )
        self.scores = self.df["MOS_zscore"].to_numpy()

        self.transform = transform

    def __getitem__(self, index: int) -> Tuple[torch.Tensor, None, float]:
        x_path = self.root / self.df.at[index, "image_name"]
        score = self.scores[index]

        x = imread(x_path)
        y = torch.rand(1)

        if self.transform is not None:
            x = self.transform(x)
            return x, y, score

        x = torch.from_numpy(x).float()
        x = x.permute(2, 0, 1)
        return x, y, score

    def __len__(self) -> int:
        return len(self.df)


class LIVEitW(KonIQ10k):
    r"""A class to evaluate on the train/test/train+test the LIVE-in-the-Wild dataset.
    https://live.ece.utexas.edu/research/ChallengeDB/index.html

    Note that the class is callable. The values are returned as a result of calling the __getitem__ method.

    WARNING: This dataset contains images with different spatial resolutions.
    Hence, inference with bs > 1 may cause runtime errors.

    Args:
        root: Root directory path.
        transform: Something that pre-processes dataset images.
        subset: Which part of the dataset to use.
            Options: "train", "test", "all". "all" means train + test.
    Returns:
        x: image with some kind of distortion in [0, 1] range.
        y: dummy variable, used for compatibility with FR datasets.
        score: MOS score of the image quality.
    """
    _filename_names = "AllImages_release.mat"
    _filename_mos = "AllMOS_release.mat"

    def __init__(self, root: Path, transform: Optional[Callable] = None, subset: str = 'test') -> None:
        supported_subsets = ["train", "test", "all"]
        assert subset in supported_subsets, f"Unknown subset [{subset}], choose one of {supported_subsets}."
        assert root.exists(), "You need to download LIVEitW dataset first."

        labels_folder = "Data"
        names = loadmat(root / labels_folder / self._filename_names)
        mos = loadmat(root / labels_folder / self._filename_mos)

        images_folder = "Images"
        n_train_images = 7  # There are only 7 images in the train set that are placed in different folder.
        train_paths = [root / images_folder / "trainingImages" / n[0][0]
                       for n in names["AllImages_release"]][:n_train_images]
        test_paths = [root / images_folder / n[0][0] for n in names["AllImages_release"]][n_train_images:]
        scores = mos["AllMOS_release"][0]

        if subset == "train":
            self.df = pd.DataFrame().from_dict({"image_name": train_paths})
            self.scores = scores[:n_train_images]
        elif subset == "test":
            self.df = pd.DataFrame().from_dict({"image_name": test_paths})
            self.scores = scores[n_train_images:]
        else:
            self.df = pd.DataFrame().from_dict({"image_name": train_paths + test_paths})
            self.scores = scores

        self.root = root
        self.transform = transform


DATASETS = {
    "tid2013": TID2013,
    "kadid10k": KADID10k,
    "pipal": PIPAL,
    "koniq10k": KonIQ10k,
    "liveitw": LIVEitW,
}


def eval_metric(loader: DataLoader, metric: Metric, device: str, feature_extractor: str) \
        -> Tuple[torch.Tensor, torch.Tensor]:
    r"""Evaluate metric on a given dataset.

    Args:
        loader: PyTorch dataloader that returns batch of distorted images, reference images and scores.
        metric: General metric that satisfies the Metric interface.
        device: Computation device.
        feature_extractor: name of the neural network to be used to extract features from images
    Returns:
        gt_scores: A tensor of ground truth values.
        metric_scores: A tensor of predicted values.
    """
    assert isinstance(loader, DataLoader), "Expect loader to be DataLoader class"
    assert isinstance(metric, Metric), f"Expected metric to be an instance of Metric, got {type(metric)} instead!"

    gt_scores = []
    metric_scores = []
    compute_function = determine_compute_function(metric_category=metric.category)

    for distorted_images, reference_images, scores in tqdm.tqdm(loader, ncols=50):
        distorted_images, reference_images = distorted_images.to(device), reference_images.to(device)
        gt_scores.append(scores.cpu())

        metric_score: torch.Tensor = \
            compute_function(metric.functor, distorted_images, reference_images, device, feature_extractor)

        if metric_score.dim() == 0:
            metric_score = metric_score.unsqueeze(0)

        metric_scores.append(metric_score.cpu())

    return torch.cat(gt_scores), torch.cat(metric_scores)


def determine_compute_function(metric_category: str) -> Callable:
    return {
        "FR": compute_full_reference,
        "NR": compute_no_reference,
        "DB": compute_distribution_based,
    }[metric_category]


def get_feature_extractor(feature_extractor_name: str, device: str) -> nn.Module:
    r"""A factory to initialize feature extractor from its name."""
    if feature_extractor_name == "vgg16":
        return torchvision.models.vgg16(pretrained=True, progress=True).features.to(device)
    elif feature_extractor_name == "vgg19":
        return torchvision.models.vgg19(pretrained=True, progress=True).features.to(device)
    elif feature_extractor_name == "inception":
        return piq.feature_extractors.InceptionV3(
            resize_input=False, use_fid_inception=True, normalize_input=True).to(device)
    else:
        raise ValueError(f"Wrong feature extractor name {feature_extractor_name}")


def compute_full_reference(
    metric_functor: Callable,
    distorted_images: torch.Tensor,
    reference_images: torch.Tensor,
    _,
    __,
) -> torch.Tensor:
    return metric_functor(distorted_images, reference_images).cpu()


def compute_no_reference(
        metric_functor: Callable,
        distorted_images: torch.Tensor,
        _,
        device: str,
        ___) -> torch.Tensor:
    if not isinstance(metric_functor, types.FunctionType):
        metric_functor = metric_functor.to(device)

    return metric_functor(distorted_images).cpu()


def extract_features(
    distorted_patches: torch.Tensor,
    feature_extractor: nn.Module,
    feature_extractor_name: str,
    reference_patches: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    distorted_features, reference_features = [], []
    with torch.no_grad():
        if feature_extractor_name == "inception":
            reference_features.append(feature_extractor(reference_patches)[0].squeeze())
            distorted_features.append(feature_extractor(distorted_patches)[0].squeeze())
        elif feature_extractor_name in ["vgg16", "vgg19"]:
            reference_features.append(torch.nn.functional.avg_pool2d(feature_extractor(reference_patches), 3).squeeze())
            distorted_features.append(torch.nn.functional.avg_pool2d(feature_extractor(distorted_patches), 3).squeeze())
        else:
            raise ValueError(f'Unknown feature extractor {feature_extractor_name} is selected. '
                             f'Please choose on of supported feature extractors: [inception, vgg16, vgg19]')

    distorted_features = torch.cat(distorted_features, dim=0)
    reference_features = torch.cat(reference_features, dim=0)

    return distorted_features, reference_features


def normalize_tensor(tensor: torch.Tensor) -> torch.Tensor:
    r"""Map tensor values to [0, 1]"""
    return (tensor - tensor.min()) / (tensor.max() - tensor.min())


def compute_distribution_based(
    metric_functor: Callable,
    distorted_images: torch.Tensor,
    reference_images: torch.Tensor,
    device: str,
    feature_extractor_name: str,
) -> torch.Tensor:
    feature_extractor = get_feature_extractor(
        feature_extractor_name=feature_extractor_name, device=device
    )

    if feature_extractor_name == "inception":
        distorted_images = normalize_tensor(distorted_images)
        reference_images = normalize_tensor(reference_images)

    # Create patches
    distorted_patches = crop_patches(distorted_images, size=96, stride=32)
    reference_patches = crop_patches(reference_images, size=96, stride=32)

    # Extract features from distorted images
    distorted_patches = distorted_patches.view(-1, *distorted_patches.shape[-3:])
    reference_patches = reference_patches.view(-1, *reference_patches.shape[-3:])

    distorted_features, reference_features = extract_features(
        distorted_patches, feature_extractor, feature_extractor_name, reference_patches
    )

    return metric_functor(distorted_features, reference_features).cpu()


def crop_patches(images: torch.Tensor, size: int = 64, stride: int = 32) -> torch.Tensor:
    r"""Crop input images into smaller patches.

    Args:
        images: Tensor of images with shape (batch x 3 x H x W)
        size: size of a square patch
        stride: Step between patches
    Returns:
        A tensor on cropped patches of shape (-1, 3, size, size)
    """
    patches = images.data.unfold(1, 3, 3).unfold(2, size, stride).unfold(3, size, stride)
    patches = patches.reshape(-1, 3, size, size)
    return patches


def main(
    dataset_name: str,
    path: Path,
    metrics: List[str],
    batch_size: int,
    device: str,
    feature_extractor: str,
) -> None:
    # Init dataset and dataloader
    dataset = DATASETS[dataset_name](root=path)
    loader = DataLoader(dataset, batch_size=batch_size, num_workers=4)

    # If category of metrics is selected instead of a list of metrics, take all metrics from this category
    if metrics[0] in METRIC_CATEGORIES:
        metrics = METRIC_CATEGORIES[metrics[0]]

    if metrics[0] == "all":
        metrics = list(chain(*METRIC_CATEGORIES.values()))

    for metric_name in metrics:
        metric: Metric = METRICS[metric_name]
        gt_scores, metric_scores = eval_metric(loader, metric, device=device, feature_extractor=feature_extractor)
        gt_scores, metric_scores = gt_scores.numpy(), metric_scores.numpy()
        print(f"{metric_name}: SRCC {abs(spearmanr(gt_scores, metric_scores)[0]):0.3f}",
              f"KRCC {abs(kendalltau(gt_scores, metric_scores)[0]):0.3f}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Benchmark PIQ metrics")

    parser.add_argument("--dataset", type=str, help="Dataset name", choices=list(DATASETS.keys()))
    parser.add_argument("--path", type=Path, help="Path to dataset")
    parser.add_argument('--metrics', nargs='+', default=[], help='Metrics to benchmark',
                        choices=list(METRICS.keys()) + list(METRIC_CATEGORIES.keys()) + ['all'])
    parser.add_argument('--batch_size', type=int, default=1, help='Batch size')
    device_choice = ['cpu', 'cuda'] + [f'cuda:{i}' for i in range(torch.cuda.device_count())]
    parser.add_argument('--device', default='cuda', choices=device_choice, help='Computation device')
    parser.add_argument('--feature_extractor', default='inception', choices=['inception', 'vgg16', 'vgg19'],
                        help='Select a feature extractor. For distribution-based metrics only')

    args = parser.parse_args()
    print(f"Parameters used for benchmark: {args}")

    main(
        dataset_name=args.dataset,
        path=args.path,
        metrics=args.metrics,
        batch_size=args.batch_size,
        device=args.device,
        feature_extractor=args.feature_extractor
    )