reconstruction.py

import os
import sys
import math
import torch
import numpy as np
import cv2

def normalize(x):
    return (x - x.min()) / (x.max() - x.min())

class PSNR:
    """Peak Signal to Noise Ratio
    img1 and img2 have range [0, 255]"""

    def __init__(self):
        self.name = "PSNR"

    @staticmethod
    def __call__(img1, img2):
        img1 = normalize(img1)*255
        img2 = normalize(img2)*255

        mse = torch.mean((img1 - img2) ** 2)
        return 20 * torch.log10(255.0 / torch.sqrt(mse))


class SSIM:
    """Structure Similarity
    img1, img2: [0, 255]"""

    def __init__(self):
        self.name = "SSIM"

    @staticmethod
    def __call__(img1, img2, reduction='mean'):
        img1 = normalize(img1)*255
        img2 = normalize(img2)*255

        if 'torch' in str(type(img1)):
            img1 = img1.detach().cpu().numpy()
            img2 = img2.detach().cpu().numpy()

        if not img1.shape == img2.shape:
            raise ValueError("Input images must have the same dimensions.")
        if img1.ndim == 2:  # Grey or Y-channel image
            return SSIM._ssim(img1, img2)
        elif img1.ndim == 3:
            return SSIM._ssim3(img1, img2)
        elif img1.ndim == 4:
            return SSIM._ssim4(img1, img2, reduction=reduction)
        else:
            raise ValueError("Wrong input image dimensions.")

    @staticmethod
    def _ssim(img1, img2):
        C1 = (0.01 * 255) ** 2
        C2 = (0.03 * 255) ** 2

        img1 = img1.astype(np.float64)
        img2 = img2.astype(np.float64)
        kernel = cv2.getGaussianKernel(11, 1.5)
        window = np.outer(kernel, kernel.transpose())

        mu1 = cv2.filter2D(img1, -1, window)[5:-5, 5:-5]  # valid
        mu2 = cv2.filter2D(img2, -1, window)[5:-5, 5:-5]
        mu1_sq = mu1 ** 2
        mu2_sq = mu2 ** 2
        mu1_mu2 = mu1 * mu2
        sigma1_sq = cv2.filter2D(img1 ** 2, -1, window)[5:-5, 5:-5] - mu1_sq
        sigma2_sq = cv2.filter2D(img2 ** 2, -1, window)[5:-5, 5:-5] - mu2_sq
        sigma12 = cv2.filter2D(img1 * img2, -1, window)[5:-5, 5:-5] - mu1_mu2

        ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / (
            (mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2)
        )
        return ssim_map.mean()

    @staticmethod
    def _ssim3(img1, img2):
        if img1.shape[0] == 3:
            ssims = []
            for i in range(3):
                ssims.append(SSIM._ssim(img1[i], img2[i]))
            return np.array(ssims).mean()
        elif img1.shape[0] == 1:
            return SSIM._ssim(np.squeeze(img1), np.squeeze(img2))

    @staticmethod
    def _ssim4(b_img1, b_img2, reduction='mean'):
        batch_size = b_img1.shape[0]
        ssims = np.zeros(batch_size)

        for i in range(batch_size):
            ssims[i] = SSIM._ssim3(b_img1[i], b_img2[i])

        if reduction == 'none':
            return ssims

        ssim = ssims.mean()
        return ssim