マシュマロンのブログ

おはようございます。こんにちは。こんばんは。

Python de PSNRのお勉強

PSNRの仕組みを勉強するためにPythonでコードを書きました。
アルゴリズムは完全に理解しましたが(本当に?)、使いどころまだよくわかってないので要勉強。
jpeg圧縮のコード(=imgEncodeDecode())も含んでいます。
※実はOpenCVにcv2.psnr()がある



import cv2
import math


def imgEncodeDecode(in_imgs, quality, ch=1):

    encode_param = [int(cv2.IMWRITE_JPEG_QUALITY), quality]

    result, encimg = cv2.imencode('.jpg', in_imgs, encode_param)
    if False == result:
        print('could not encode image!')
        exit()

    decimg = cv2.imdecode(encimg, ch)
    
    return decimg

#↑参考:https://qiita.com/ka10ryu1/items/5fed6b4c8f29163d0d65

    
def PSNR(ori_img, dec_img):
    height, width, BPP = ori_img.shape
    N = height*width
    
    sumR = 0
    sumG = 0
    sumB = 0
    
    for i in range(height):
        for j in range(width):
            sumR += (int(ori_img[i][j][2]) - int(dec_img[i][j][2]))**2
            sumG += (int(ori_img[i][j][1]) - int(dec_img[i][j][1]))**2
            sumB += (int(ori_img[i][j][0]) - int(dec_img[i][j][0]))**2
            
    MSE = (sumR + sumG + sumB)/(3*N)
    
    if MSE == 0:
        ans = "original"
    else:
        ans =10*math.log10((255**2)/MSE)
    return ans


    
ori_img = cv2.imread("XXX.jpg")
dec_img = imgEncodeDecode(ori_img, 50)
cv2.imwrite('Q50.jpg', dec_img)
print(PSNR(ori_img, dec_img))