major

main.py

from math import log10, sqrt
import cv2
import numpy as np
from skimage.metrics import structural_similarity
from piqa import ssim
from PIL import Image
from tqdm import tqdm
from pytorch_msssim import ssim, ms_ssim, SSIM, MS_SSIM
import torchvision
from brisque import BRISQUE

topil = torchvision.transforms.ToPILImage()
totensor = torchvision.transforms.ToTensor()


###################### PSNR

def PSNR(original, compare):
    mse = np.mean((original - compare) ** 2)
    if mse == 0:  # MSE is zero means no noise is present in the signal .
        # Therefore PSNR have no importance.
        return 100
    max_pixel = 255.0
    psnr = 20 * log10(max_pixel / sqrt(mse))
    print(f"PSNR: {psnr} dB")


###################### SSIM

def SSIM(original, compare):
    # Convert images to grayscale
    before_gray = cv2.cvtColor(original, cv2.COLOR_BGR2GRAY)
    after_gray = cv2.cvtColor(compare, cv2.COLOR_BGR2GRAY)

    # Compute SSIM between two images
    (score, diff) = structural_similarity(before_gray, after_gray, full=True)

    # The diff image contains the actual image differences between the two images
    # and is represented as a floating point data type in the range [0,1]
    # so we must convert the array to 8-bit unsigned integers in the range
    # [0,255] before we can use it with OpenCV
    diff = (diff * 255).astype("uint8")

    print("SSIM: {}".format(score))


###################### MS-SSIM

def MSSSIM(original, compare):
    it1 = totensor(original).unsqueeze(0)
    it2 = totensor(compare).unsqueeze(0)

    msssim = ssim.msssim(it1, it2).squeeze(0)

    print(f"MS-SSIM: {msssim}")


###################### MAIN

def main():
    for x in range(1, 6):
        k = "0" + str(x) + "_outdoor_hazy.jpg"
        obj = BRISQUE(k, url=False)
        print("BRISQUE: {}".format(obj.score()))

    for x in range(1, 46):
        if x < 10:
            original = cv2.imread("GT/" + "0" + str(x) + "_outdoor_GT.jpg")
            compare = cv2.imread("hazy/" + "0" + str(x) + "_outdoor_hazy.jpg")
        else:
            original = cv2.imread("GT/" + str(x) + "_outdoor_GT.jpg")
            compare = cv2.imread("hazy/" + str(x) + "_outdoor_hazy.jpg")
        print(f"{x} ===========")
        PSNR(original, compare)
        SSIM(original, compare)
        # MSSSIM(original, compare)


if __name__ == "__main__":
    main()