Compute libopus and libvorbis PSNR

opus_psrn.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import scipy.io.wavfile
import scipy.signal
import subprocess
import tempfile

ffmpeg = "avconv"

def psnr(a, b):
	import numpy
	import math
	
	n = min(len(a), len(b))
	mse = 0
	for i in xrange(n):
		for j in xrange(2):
			mse += (int(a[i][j]) - int(b[i][j])) ** 2
	mse = mse / float(n)
	if mse == 0:
		return None
	return 20 * math.log10(2**16 / math.sqrt(mse))

def towav(fn, forceRate=None):
	with tempfile.NamedTemporaryFile() as temp_wav:
		subprocess.Popen([ffmpeg, '-i', fn] + (['-ar', str(forceRate)] if forceRate is not None else []) + ['-f', 'wav', temp_wav.name, '-y']).communicate()
		rate,wav = scipy.io.wavfile.read(temp_wav)
		if forceRate is not None:
			assert rate == forceRate, "{} {}".format(rate, forceRate)
	return wav,rate

def test(wavfile):
	orig,rate = towav(wavfile)

	with tempfile.NamedTemporaryFile() as temp:
		subprocess.Popen([ffmpeg, '-i', wavfile, '-c:a', 'libopus', '-b:a', '64k', '-f', 'ogg', temp.name, '-y']).communicate()
		opus_wav = towav(temp.name, rate)[0]

	with tempfile.NamedTemporaryFile() as temp:
		subprocess.Popen([ffmpeg, '-i', wavfile, '-c:a', 'libvorbis', '-aq', '4', '-f', 'ogg', temp.name, '-y']).communicate()
		vorbis_wav = towav(temp.name, rate)[0]

	opus_psnr = psnr(orig, opus_wav)
	vorbis_psnr = psnr(orig, vorbis_wav)

	print "RESULTS: libopus PSNR {}, libvorbis PSNR {}".format(opus_psnr, vorbis_psnr)

if __name__ == '__main__':
	import argparse
	parser = argparse.ArgumentParser()
	parser.add_argument('wavfile', type=str)
	parser.add_argument('--ffmpeg', action="store_true")
	opts = parser.parse_args()

	if opts.ffmpeg:
		ffmpeg = "ffmpeg"

	test(opts.wavfile)