两个图像的 PSNR/MSE 计算

Question

我想编写一个函数，它获取两个图像参考和编码，并评估每个组件（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 秒。

Answer 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) 索引的 dict。

---

编辑

我尝试使用 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)的numpy数组。我认为您不需要其中的原始 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) 将使像素通道成为“外部”（变化最慢的）索引。