MSE partitiioned by distance from edge

edge_psnr.py

#!/bin/env python3

import array
import sys
import y4m

ref_frame = None


def frame_arrays(frame):
    w = int(frame.headers['W'])
    h = int(frame.headers['H'])
    wc = (w + 1) // 2
    hc = (h + 1) // 2
    y = array.array('B', frame.buffer[:w * h])
    u = array.array('B', frame.buffer[w * h:w * h + wc * hc])
    v = array.array('B', frame.buffer[w * h + wc * hc:])
    return w, h, y, u, v


def process_ref(frame):
    global ref_frame
    ref_frame = frame_arrays(frame)


rec_frame = None


def process_rec(frame):
    global rec_frame
    rec_frame = frame_arrays(frame)


if __name__ == '__main__':
    parser = y4m.Reader(process_ref)
    with open(sys.argv[1], 'rb') as f:
        while True:
            data = f.read(4096)
            if not data:
                break
            parser.decode(data)
    parser = y4m.Reader(process_rec)
    with open(sys.argv[2], 'rb') as f:
        while True:
            data = f.read(4096)
            if not data:
                break
            parser.decode(data)
    w, h, y_ref, u_ref, v_ref = ref_frame
    _, _, y_rec, u_rec, v_rec = rec_frame

    count = array.array('L', [0] * 64)
    sse = array.array('L', [0] * 64)

    for y in range(0, h):
        for x in range(0, w):
            xy = y * w + x
            i = min(64, w - 1 - x, h - 1 - y)
            if i == 64: continue
            count[i] += 1
            sse[i] += (y_rec[xy] - y_ref[xy])**2
    wc = (w + 1) // 2
    hc = (h + 1) // 2
    for y in range(0, hc):
        for x in range(0, wc):
            xy = y * wc + x
            i = min(64, w - 1 - x * 2, h - 1 - y * 2)
            if i == 64: continue
            count[i] += 2
            sse[i] += (u_rec[xy] - u_ref[xy])**2
            sse[i] += (v_rec[xy] - v_ref[xy])**2
            i = min(64, w - 2 - x * 2, h - 2 - y * 2)
            if i == 64: continue
            count[i] += 2
            sse[i] += (u_rec[xy] - u_ref[xy])**2
            sse[i] += (v_rec[xy] - v_ref[xy])**2
    scale = sum(count) * 1. / sum(sse)
    for i in range(0, 64):
        print("%d %d %d" % (i, sse[i] * scale, count[i]))