两幅图像的PSNR / MSE计算

时间:2018-10-26 09:28:48

标签: python image rgb mse

我想编写一个函数,该函数获取两个图像参考并进行编码,并针对每个分量(R,G,B,Y,Cb,Cr)评估(R)MSE和PSNR。为此,我提取所有分量,然后转换RGB-> YCbCr。我想在不使用内置函数的情况下计算(R)MSE和PSNR。

import os, sys, subprocess, csv, datetime
from PIL import Image

############ Functions Definitions ############

# Extracts the values of the R, G, B components of each pixel in the input file and calculates the Y, Cb, Cr components returning a dictionary having a key tuple with the coordinates of
 the pixes and values the values of each R, G, B, Y, Cb, Cr components
def rgb_calc(ref_file):
  img = Image.open(ref_file)
  width, height = img.size
  print(width)
  print(height)
  rgb_dict = {}
  for x in range (width):
    for y in range(height):
      r, g, b = img.load()[x, y]
      lum = 0.299 * r + 0.587 * g + 0.114 * b
      cb = 128 - 0.168736 * r - 0.331264 * g + 0.5 * b
      cr = 128 + 0.5 * r - 0.418688 * g - 0.081312 * b
      print("X {} Y {} R {} G {} B {} Y {} Cb {} Cr {}".format(x, y, r, g, b, lum, cb, cr))
      rgb_dict[(x, y)] = (r, g, b, lum, cb, cr)
  return rgb_dict

############ MAIN FUNCTION ############

r_img = sys.argv[1]
p_img = sys.argv[2]

ref_img = Image.open(r_img)
proc_img = Image.open(p_img)

resolution_ref = ref_img.size
resolution_proc = proc_img.size

if resolution_ref == resolution_proc:
  ycbcr_ref = rgb_calc(r_img)
  ycbcr_proc = rgb_calc(proc_img)
else:
  exit(0)

我想编写一个新函数,并最终输出每个分量的平均PSNR和整个图像的平均值。

是否可以加快我的流程?

当前,img.load()每8Mpx图像大约需要10-11秒,而字典的创建又需要6秒。因此,仅提取这些值并创建两个字典已经花费了32秒。

1 个答案:

答案 0 :(得分:1)

首先,在循环之外执行img.load()

def rgb_calc(ref_file):
  img = Image.open(ref_file)
  width, height = img.size
  print(width)
  print(height)
  rgb_dict = {}
  rgb = img.load()
  for x in range(width):
    for y in range(height):
      r, g, b = rgb[x, y]
      lum = 0.299 * r + 0.587 * g + 0.114 * b
      cb = 128 - 0.168736 * r - 0.331264 * g + 0.5 * b
      cr = 128 + 0.5 * r - 0.418688 * g - 0.081312 * b
      rgb_dict[(x, y)] = (r, g, b, lum, cb, cr)
  return rgb_dict

但这只是开始。接下来我要做的一件事(但我不是专家!)是使用numpy数组,而不是用(x,y)索引的字典。


编辑

我试图使用numpy ndarray(N维数组)加快处理速度,但是被卡住了,所以问了一个具体问题,并得到了解决的答案(加快了15倍!): numpy.ndarray with shape (height, width, n) from n values per Image pixel

在这里,它可以根据您的需要进行调整,并固定了一些原始代码:

import numpy as np
from PIL import Image

def get_rgbycbcr(img: Image.Image):
    R, G, B = np.array(img).transpose(2, 0, 1)[:3]  # ignore alpha if present
    Y = 0.299 * R + 0.587 * G + 0.114 * B
    Cb = 128 - 0.168736 * R - 0.331264 * G + 0.5 * B
    Cr = 128 + 0.5 * R - 0.418688 * G - 0.081312 * B
    return np.array([R, G, B, Y, Cb, Cr], dtype=float).transpose(2, 1, 0)

r_img = sys.argv[1]
p_img = sys.argv[2]

ref_img  = Image.open(r_img)
proc_img = Image.open(p_img)

resolution_ref  = ref_img.size
resolution_proc = proc_img.size

if resolution_ref == resolution_proc:
    ycbcr_ref  = get_ycbcr(ref_img) 
    ycbcr_proc = get_ycbcr(proc_img)
else:
    exit(0)

现在剩下的是一个形状为(width, height, 6)的数字数组。我认为您不需要那里的原始RGB数据(您可以随时从图像中获取它),以防万一,您可以将代码更改为6到3。您可以像这样对ycbcr_ref进行索引:ycbcr_ref[x, y]并获得长度为6的列表,其中包含与存储在字典中的元组中相同的数据。但是您可以提取切片,特别是沿着长度为6的“轴”(numpy术语)提取切片,并对切片进行操作,例如

y_mean = ycbcr_ref[:, :, 3].mean()

学习如何使用numpy绝对值得!

我将为您提供一个详细信息:除非另行说明,否则numpy会以变化最快的索引(AKA轴)为首,最后存储最快的变化来存储数据。由于图像是按行存储的,因此除非您执行transpose(),否则读入numpy的图像必须像arr[y, x]一样进行索引。移置将使轴移动。在您的情况下,您有3个轴,编号为0、1、2。例如,.transpose(1, 0, 2)将交换x和y,而.transpose(2, 0, 1)将使像素通道成为“外部”(变化最小的)索引。