转自github上的生成对抗网络超分辨率重建

# Copyright 2022 Dakewe Biotech Corporation. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import collections.abc
import math
import typing
import warnings
from itertools import repeat
from typing import Any

import cv2
import numpy as np
import torch
from numpy import ndarray
from scipy.io import loadmat
from scipy.ndimage.filters import convolve
from scipy.special import gamma
from torch import nn
from torch.nn import functional as F

from imgproc import image_resize, expand_y, bgr_to_ycbcr, rgb_to_ycbcr_torch

__all__ = [
    "psnr", "ssim", "niqe",
    "PSNR", "SSIM", "NIQE",
]

_I = typing.Optional[int]
_D = typing.Optional[torch.dtype]


# The following is the implementation of IQA method in Python, using CPU as processing device
def _check_image(raw_image: np.ndarray, dst_image: np.ndarray):
    """Check whether the size and type of the two images are the same

    Args:
        raw_image (np.ndarray): image data to be compared, BGR format, data range [0, 255]
        dst_image (np.ndarray): reference image data, BGR format, data range [0, 255]

    """
    # check image scale
    assert raw_image.shape == dst_image.shape, \
        f"Supplied images have different sizes {str(raw_image.shape)} and {str(dst_image.shape)}"

    # check image type
    if raw_image.dtype != dst_image.dtype:
        warnings.warn(f"Supplied images have different dtypes{str(raw_image.shape)} and {str(dst_image.shape)}")


def psnr(raw_image: np.ndarray, dst_image: np.ndarray, crop_border: int, only_test_y_channel: bool) -> float:
    """Python implements PSNR (Peak Signal-to-Noise Ratio, peak signal-to-noise ratio) function

    Args:
        raw_image (np.ndarray): image data to be compared, BGR format, data range [0, 255]
        dst_image (np.ndarray): reference image data, BGR format, data range [0, 255]
        crop_border (int): crop border a few pixels
        only_test_y_channel (bool): Whether to test only the Y channel of the image.

    Returns:
        psnr_metrics (np.float64): PSNR metrics

    """
    # Check if two images are similar in scale and type
    _check_image(raw_image, dst_image)

    # crop border pixels
    if crop_border > 0:
        raw_image = raw_image[crop_border:-crop_border, crop_border:-crop_border, ...]
        dst_image = dst_image[crop_border:-crop_border, crop_border:-crop_border, ...]

    # If you only test the Y channel, you need to extract the Y channel data of the YCbCr channel data separately
    if only_test_y_channel:
        raw_image = expand_y(raw_image)
        dst_image = expand_y(dst_image)

    # Convert data type to numpy.float64 bit
    raw_image = raw_image.astype(np.float64)
    dst_image = dst_image.astype(np.float64)

    psnr_metrics = 10 * np.log10((255.0 ** 2) / np.mean((raw_image - dst_image) ** 2) + 1e-8)

    return psnr_metrics


def _ssim(raw_image: np.ndarray, dst_image: np.ndarray) -> float:
    """Python implements the SSIM (Structural Similarity) function, which only calculates single-channel data

    Args:
        raw_image (np.ndarray): The image data to be compared, in BGR format, the data range is [0, 255]
        dst_image (np.ndarray): reference image data, BGR format, data range is [0, 255]

    Returns:
        ssim_metrics (float): SSIM metrics for single channel

    """
    c1 = (0.01 * 255.0) ** 2
    c2 = (0.03 * 255.0) ** 2

    kernel = cv2.getGaussianKernel(11, 1.5)
    kernel_window = np.outer(kernel, kernel.transpose())

    raw_mean = cv2.filter2D(raw_image, -1, kernel_window)[5:-5, 5:-5]
    dst_mean = cv2.filter2D(dst_image, -1, kernel_window)[5:-5, 5:-5]
    raw_mean_square = raw_mean ** 2
    dst_mean_square = dst_mean ** 2
    raw_dst_mean = raw_mean * dst_mean
    raw_variance = cv2.filter2D(raw_image ** 2, -1, kernel_window)[5:-5, 5:-5] - raw_mean_square
    dst_variance = cv2.filter2D(dst_image ** 2, -1, kernel_window)[5:-5, 5:-5] - dst_mean_square
    raw_dst_covariance = cv2.filter2D(raw_image * dst_image, -1, kernel_window)[5:-5, 5:-5] - raw_dst_mean

    ssim_molecular = (2 * raw_dst_mean + c1) * (2 * raw_dst_covariance + c2)
    ssim_denominator = (raw_mean_square + dst_mean_square + c1) * (raw_variance + dst_variance + c2)

    ssim_metrics = ssim_molecular / ssim_denominator
    ssim_metrics = float(np.mean(ssim_metrics))

    return ssim_metrics


def ssim(raw_image: np.ndarray, dst_image: np.ndarray, crop_border: int, only_test_y_channel: bool) -> float:
    """Python implements the SSIM (Structural Similarity) function, which calculates single/multi-channel data

    Args:
        raw_image (np.ndarray): The image data to be compared, in BGR format, the data range is [0, 255]
        dst_image (np.ndarray): reference image data, BGR format, data range is [0, 255]
        crop_border (int): crop border a few pixels
        only_test_y_channel (bool): Whether to test only the Y channel of the image

    Returns:
        ssim_metrics (float): SSIM metrics for single channel

    """
    # Check if two images are similar in scale and type
    _check_image(raw_image, dst_image)

    # crop border pixels
    if crop_border > 0:
        raw_image = raw_image[crop_border:-crop_border, crop_border:-crop_border, ...]
        dst_image = dst_image[crop_border:-crop_border, crop_border:-crop_border, ...]

    # If you only test the Y channel, you need to extract the Y channel data of the YCbCr channel data separately
    if only_test_y_channel:
        raw_image = expand_y(raw_image)
        dst_image = expand_y(dst_image)

    # Convert data type to numpy.float64 bit
    raw_image = raw_image.astype(np.float64)
    dst_image = dst_image.astype(np.float64)

    channels_ssim_metrics = []
    for channel in range(raw_image.shape[2]):
        ssim_metrics = _ssim(raw_image[..., channel], dst_image[..., channel])
        channels_ssim_metrics.append(ssim_metrics)
    ssim_metrics = np.mean(np.asarray(channels_ssim_metrics))

    return float(ssim_metrics)


def _estimate_aggd_parameters(vector: np.ndarray) -> [np.ndarray, float, float]:
    """Python implements the NIQE (Natural Image Quality Evaluator) function,
    This function is used to estimate an asymmetric generalized Gaussian distribution

    Reference papers:
        `Estimation of shape parameter for generalized Gaussian distributions in subband decompositions of video`

    Args:
        vector (np.ndarray): data vector

    Returns:
        aggd_parameters (np.ndarray): asymmetric generalized Gaussian distribution
        left_beta (float): symmetric left data vector variance mean product
        right_beta (float): symmetric right side data vector variance mean product

    """
    # The following is obtained according to the formula and the method provided in the paper on WIki encyclopedia
    vector = vector.flatten()
    gam = np.arange(0.2, 10.001, 0.001)  # len = 9801
    gam_reciprocal = np.reciprocal(gam)
    r_gam = np.square(gamma(gam_reciprocal * 2)) / (gamma(gam_reciprocal) * gamma(gam_reciprocal * 3))

    left_std = np.sqrt(np.mean(vector[vector < 0] ** 2))
    right_std = np.sqrt(np.mean(vector[vector > 0] ** 2))
    gamma_hat = left_std / right_std
    rhat = (np.mean(np.abs(vector))) ** 2 / np.mean(vector ** 2)
    rhat_norm = (rhat * (gamma_hat ** 3 + 1) * (gamma_hat + 1)) / ((gamma_hat ** 2 + 1) ** 2)
    array_position = np.argmin((r_gam - rhat_norm) ** 2)

    aggd_parameters = gam[array_position]
    left_beta = left_std * np.sqrt(gamma(1 / aggd_parameters) / gamma(3 / aggd_parameters))
    right_beta = right_std * np.sqrt(gamma(1 / aggd_parameters) / gamma(3 / aggd_parameters))

    return aggd_parameters, left_beta, right_beta


def _get_mscn_feature(image: np.ndarray) -> Any:
    """Python implements the NIQE (Natural Image Quality Evaluator) function,
    This function is used to calculate the MSCN feature map

    Reference papers:
        `Estimation of shape parameter for generalized Gaussian distributions in subband decompositions of video`

    Args:
        image (np.ndarray): Grayscale image of MSCN feature to be calculated, BGR format, data range is [0, 255]

    Returns:
        mscn_feature (Any): MSCN feature map of the image

    """
    mscn_feature = []
    # Calculate the asymmetric generalized Gaussian distribution
    aggd_parameters, left_beta, right_beta = _estimate_aggd_parameters(image)
    mscn_feature.extend([aggd_parameters, (left_beta + right_beta) / 2])

    shifts = [[0, 1], [1, 0], [1, 1], [1, -1]]
    for i in range(len(shifts)):
        shifted_block = np.roll(image, shifts[i], axis=(0, 1))
        # Calculate the asymmetric generalized Gaussian distribution
        aggd_parameters, left_beta, right_beta = _estimate_aggd_parameters(image * shifted_block)
        mean = (right_beta - left_beta) * (gamma(2 / aggd_parameters) / gamma(1 / aggd_parameters))
        mscn_feature.extend([aggd_parameters, mean, left_beta, right_beta])

    return mscn_feature


def _fit_mscn_ipac(image: np.ndarray,
                   mu_pris_param: np.ndarray,
                   cov_pris_param: np.ndarray,
                   gaussian_window: np.ndarray,
                   block_size_height: int,
                   block_size_width: int) -> float:
    """Python implements the NIQE (Natural Image Quality Evaluator) function,
    This function is used to fit the inner product of adjacent coefficients of MSCN

    Reference papers:
        `Estimation of shape parameter for generalized Gaussian distributions in subband decompositions of video`

    Args:
        image (np.ndarray): The image data of the NIQE to be tested, in BGR format, the data range is [0, 255]
        mu_pris_param (np.ndarray): Mean of predefined multivariate Gaussians, model computed on original dataset.
        cov_pris_param (np.ndarray): Covariance of predefined multivariate Gaussian model computed on original dataset.
        gaussian_window (np.ndarray): 7x7 Gaussian window for smoothing the image
        block_size_height (int): the height of the block into which the image is divided
        block_size_