Расчет PSNR / MSE для двух изображений

Question

Я хочу написать функцию, которая получает два изображения, эталонные и кодированные, и оценивает (R) MSE и PSNR для каждого компонента (R, G, B, Y, Cb, Cr).Для этого я извлекаю все компоненты, а затем конвертирую 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() занимает около 10-11 секунд на 8-мегапиксельное изображение и создание словарядополнительные 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

Но это только начало.Следующее, что я хотел бы сделать (но я не эксперт!), Это использовать пустой массив вместо dict, проиндексированного (x, y).

---

EDIT

Я пыталсячтобы ускорить процесс, используя numpy ndarray (N-мерный массив), но застрял, поэтому задал конкретный вопрос и получил решающий ответ (ускорение × 15!): numpy.ndarray с формой (высота), ширина, n) из n значений на пиксель изображения

Вот оно, адаптированное к вашим потребностям, с фиксированной детализацией вашего исходного кода:

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)

Что выосталось с теперь это numpy массив формы (width, height, 6).Я не думаю, что вам нужны исходные данные RGB (вы можете получить их в любое время из картинки) - вы можете изменить код, уменьшив с 6 до 3, на всякий случай.Вы можете индексировать, например, ycbcr_ref следующим образом: ycbcr_ref[x, y] и получить список длиной 6, содержащий те же данные, которые были у вас в кортежах, хранящихся в словаре.Но вы можете извлекать срезы, в частности, по этой «оси» (терминология numpy) длиной 6, и выполнять над ними операции, например

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

. Стоит научиться использовать numpy * 1024.*!

Я помогу вам с одной деталью: если вы не укажете обратное, numpy сохраняет данные с самым медленным изменяющимся индексом (ось AKA) первым и самым быстрым изменяющимся последним.Поскольку изображения хранятся по строкам, если вы не сделаете transpose(), изображение, считанное в numpy, должно быть проиндексировано как arr[y, x].Транспонирование будет перетасовывать оси.В вашем случае у вас есть 3 оси с номерами 0, 1, 2. Например, .transpose(1, 0, 2) поменяет местами x и y, а .transpose(2, 0, 1) сделает пиксельные каналы «внешним» (самым медленным изменяющимся) индексом.