predict_one_image.py

"""Inference wrapper for the CKDN image-quality-assessment model.

Typical usage::

    scorer = IQA_CKDN(checkpoint="model_best.pth.tar")
    score = scorer.predict(restored_addr="img/restored/3.png",
                           degraded_addr="img/degraded/3.png")
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import logging

import torch
from PIL import Image
from timm.data import create_transform, resolve_data_config
from torchvision import transforms

import ckdn

torch.backends.cudnn.benchmark = True


class IQA_CKDN:
    """No-reference-style IQA scoring using a restored/degraded image pair."""

    def __init__(self, checkpoint="model_best.pth.tar", device=None):
        cnf = {
            "interpolation": "",
            "mean": None,
            "model": "resnet",
            "std": None,
        }

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

        # create model
        self.model = ckdn.model()
        logging.info("Model %s created, param count: %d",
                     cnf["model"], sum(m.numel() for m in self.model.parameters()))

        state_dict = torch.load(checkpoint, map_location=self.device)["state_dict"]
        # Strip the module. prefix added by DataParallel / DDP checkpoints.
        new_state_dict = {
            k.split("module.", 1)[-1]: v for k, v in state_dict.items()
        }
        self.model.load_state_dict(new_state_dict, strict=False)
        self.model.to(self.device)
        self.model.eval()

        self.config = resolve_data_config(cnf, model=self.model)

    def _load(self, img_addr):
        with torch.no_grad():
            img = Image.open(img_addr).convert("RGB")
            transform = transforms.Compose(
                [create_transform(input_size=self.config["input_size"])])
            return transform(img).unsqueeze(0).to(self.device)

    def _pair(self, restored_addr, degraded_addr):
        return self._load(restored_addr), self._load(degraded_addr)

    def predict(self, restored_addr, degraded_addr):
        """Return the predicted quality score (higher is better) for the pair."""
        with torch.no_grad():
            rest, dist = self._pair(restored_addr, degraded_addr)
            output = self.model.forward_test(rest, dist)
            return output[:, 0].cpu().numpy()[0]

    def get_DTE_features(self, img_addr):
        """Deep texture encoder features of a single image."""
        with torch.no_grad():
            return self.model.extract_DTE_features(self._load(img_addr)).cpu().numpy()

    def get_QSE_features(self, img_addr):
        """Quality-sensitive encoder features of a single image."""
        with torch.no_grad():
            return self.model.extract_QSE_features(self._load(img_addr)).cpu().numpy()

    def get_QSE_minus_DTE_features(self, restored_addr, degraded_addr):
        """QSE minus DTE residual features for the restored/degraded pair."""
        with torch.no_grad():
            rest, dist = self._pair(restored_addr, degraded_addr)
            output = self.model.extract_QSE_after_minus_features(rest, dist)
            return output.cpu().numpy()

    def extract_last_features(self, restored_addr, degraded_addr):
        """Final feature vector (pre-regression head) for the pair."""
        with torch.no_grad():
            rest, dist = self._pair(restored_addr, degraded_addr)
            output = self.model.extract_last_features(rest, dist)
            return output.cpu().numpy()


def main():
    parser = argparse.ArgumentParser(description="CKDN IQA score for a restored image")
    parser.add_argument("--image", required=True,
                        help="path to the restored image")
    parser.add_argument("--degraded", default=None,
                        help="path to the corresponding degraded image "
                             "(default: sibling <folder>/degraded/ image)")
    parser.add_argument("--checkpoint", default="model_best.pth.tar",
                        help="path to the trained checkpoint")
    args = parser.parse_args()

    if args.degraded is None:
        parts = args.image.split("/")
        if len(parts) < 2 or parts[-2] in (".", ""):
            parser.error("--degraded is required when the image is not in a "
                         "restored/<name> sibling-folder layout")
        parts[-2] = "degraded"
        args.degraded = "/".join(parts)

    scorer = IQA_CKDN(checkpoint=args.checkpoint)
    score = scorer.predict(args.image, args.degraded)
    print(score)


if __name__ == "__main__":
    main()