from math import log10, sqrt
import cv2
import numpy as np

def PSNR(original, compressed):
    '''
    Calculates the Peak signal to noise ratio between a ground truth image and predicted image.
    
    see https://www.geeksforgeeks.org/python-peak-signal-to-noise-ratio-psnr/ 
    for reference

    Parameters
    ----------
        true image (cv2 image)
        predicted image (cv2 image)
        
    Returns
    -------
        PSNR score
    '''
    mse = np.mean((original - compressed) ** 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))
    return psnr

def colourfulnessMetric(img):
    """
    Created on Mon Nov 15 10:55:16 2021

    @author: Yucheng

    Parameters
    ----------
    img : cv2 RGB image

    Returns
    -------
    M : colourness metric
    -----------------------------
    |not colourful        | 0   |
    |slightly colorful    | 15  |
    |moderately colourful | 33  | 
    |averagely colourful  | 45  | 
    |quite colourful      | 59  |
    |highly colourful     | 82  |
    |extremely colourful  | 109 |
    -----------------------------
    
    """
    # Get RGB components
    R,G,B = cv2.split(img.astype("float"))
    
    # colourfulness metric from Hasler et al., section 7
    rg = R - G
    yb = (1/2) * (R+G) - B
    
    sigma_rgyb = np.sqrt(np.var(rg) + np.var(yb))
    mu_rgyb = np.sqrt(np.mean(rg)**2 + np.mean(yb)**2)
    
    M = sigma_rgyb + 0.3 * mu_rgyb
    
    return M

def main():
    import matplotlib.pyplot as plt
    original = cv2.imread("test_imgs/original_image.png")
    compressed = cv2.imread("test_imgs/compressed_image1.png", 1)
    value = PSNR(original, compressed)
    print(f"PSNR value is {value} dB")


    img2 = cv2.imread("rainbow.jpg") # opens as BGR
    img2 = cv2.cvtColor(img2, cv2.COLOR_BGR2RGB)

    plt.imshow(img2[:,:,:])
    plt.show()

    M = colourfulnessMetric(img2)
    print(M)
       
if __name__ == "__main__":
    main()

[((21, 10, 21, 47), 'numpy.mean', 'np.mean', ({(21, 18, 21, 46): '(original - compressed) ** 2'}, {}), '((original - compressed) ** 2)', True, 'import numpy as np\n'), ((69, 15, 69, 57), 'cv2.imread', 'cv2.imread', ({(69, 26, 69, 56): '"""test_imgs/original_image.png"""'}, {}), "('test_imgs/original_image.png')", False, 'import cv2\n'), ((70, 17, 70, 65), 'cv2.imread', 'cv2.imread', ({(70, 28, 70, 61): '"""test_imgs/compressed_image1.png"""', (70, 63, 70, 64): '1'}, {}), "('test_imgs/compressed_image1.png', 1)", False, 'import cv2\n'), ((75, 11, 75, 36), 'cv2.imread', 'cv2.imread', ({(75, 22, 75, 35): '"""rainbow.jpg"""'}, {}), "('rainbow.jpg')", False, 'import cv2\n'), ((76, 11, 76, 48), 'cv2.cvtColor', 'cv2.cvtColor', ({(76, 24, 76, 28): 'img2', (76, 30, 76, 47): 'cv2.COLOR_BGR2RGB'}, {}), '(img2, cv2.COLOR_BGR2RGB)', False, 'import cv2\n'), ((78, 4, 78, 27), 'matplotlib.pyplot.imshow', 'plt.imshow', ({(78, 15, 78, 26): 'img2[:, :, :]'}, {}), '(img2[:, :, :])', True, 'import matplotlib.pyplot as plt\n'), ((79, 4, 79, 14), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((60, 25, 60, 35), 'numpy.var', 'np.var', ({(60, 32, 60, 34): 'rg'}, {}), '(rg)', True, 'import numpy as np\n'), ((60, 38, 60, 48), 'numpy.var', 'np.var', ({(60, 45, 60, 47): 'yb'}, {}), '(yb)', True, 'import numpy as np\n'), ((26, 34, 26, 43), 'math.sqrt', 'sqrt', ({(26, 39, 26, 42): 'mse'}, {}), '(mse)', False, 'from math import log10, sqrt\n'), ((61, 22, 61, 33), 'numpy.mean', 'np.mean', ({(61, 30, 61, 32): 'rg'}, {}), '(rg)', True, 'import numpy as np\n'), ((61, 39, 61, 50), 'numpy.mean', 'np.mean', ({(61, 47, 61, 49): 'yb'}, {}), '(yb)', True, 'import numpy as np\n')]