AVIF vs JPEG

encode.py

#!/usr/bin/python3

import json
import os
import struct
import subprocess
import sys
import tempfile


def search(lo, hi, costfn, target):
    mid = (lo + hi) // 2
    v = costfn(mid)
    d = abs(v - target)
    bd, bi = d, mid
    while lo < hi:
        if d < bd: bd, bi = d, mid
        if v > target:
            hi = mid - 1
        elif v < target:
            lo = mid + 1
        else:
            lo = mid
            hi = mid
        if lo < hi:
            mid = (lo + hi) // 2
            v = costfn(mid)
            d = abs(v - target)
    if lo < mid:
        v = costfn(lo)
        d = abs(v - target)
        if d < bd: bd, bi = d, lo
    if hi > mid:
        v = costfn(hi)
        d = abs(v - target)
        if d < bd: bd, bi = d, hi
    return bi


def msssim(ref, rec):
    with tempfile.NamedTemporaryFile(suffix='.json', delete=True) as fp:
        subprocess.run([
            'av-metrics-tool',
            '--metric',
            'msssim',
            '--export-json',
            fp.name,
            ref,
            rec,
        ],
                       capture_output=True)
        payload = json.load(fp.file)
    return -payload['comparisons'][0]['msssim']['avg']


def jpg2y4m(fname):
    subprocess.run([
        'ffmpeg',
        '-y',
        '-i',
        fname,
        fname + '.y4m',
    ],
                   capture_output=True)
    return fname + '.y4m'


def y4m2png(y4m, png):
    subprocess.run([
        'ffmpeg',
        '-y',
        '-i',
        y4m,
        png,
    ], capture_output=True)
    with open(png, 'rb') as fp:
        return struct.unpack('>ii', fp.read(24)[16:24])


def png2y4m(png, y4m):
    subprocess.run([
        'ffmpeg',
        '-y',
        '-i',
        png,
        '-pix_fmt',
        'yuv420p',
        y4m,
    ],
                   capture_output=True)


def avif2y4m(fname):
    subprocess.run([
        'avifdec',
        fname,
        fname + '.y4m',
    ], capture_output=True)
    return fname + '.y4m'


def match_image(orig_y4m='test.y4m',
                test_y4m='test.y4m',
                test_png='test.png',
                test_jpg='test-%s.jpg',
                test_avif='test-%s-%s.avif'):
    def cjpeg(q):
        fname = test_jpg % q
        if not os.path.exists(fname):
            subprocess.run([
                '/opt/mozjpeg/bin/cjpeg', '-outfile', fname, '-quality',
                str(q), '-sample', '2x2', '-tune-ms-ssim', test_png
            ])
        return os.stat(fname).st_size

    def avifenc(off, q):
        fname = test_avif % (q, q + off)
        if not os.path.exists(fname):
            subprocess.run([
                'avifenc',
                '-c',
                'aom',
                '-s',
                '0',
                '-a',
                'aq-mode=1',
                '-a',
                'tune=ssim',
                '--min',
                str(q),
                '--max',
                str(q + off),
                '-o',
                fname,
                '-y',
                '420',
                '-d',
                '8',
                test_png,
            ],
                           capture_output=True)
        y4m = avif2y4m(fname)
        score = msssim(test_y4m, y4m)
        os.remove(y4m)
        print(fname, score)
        return score

    width, height = y4m2png(orig_y4m, test_png)
    png2y4m(test_png, test_y4m)
    target_size = width * height * 203 // 800  # 2.03 bpp
    print(target_size)
    jpeg_q = search(50, 100, cjpeg, target_size)
    print(jpeg_q)
    jpeg_msssim = msssim(test_y4m, jpg2y4m(test_jpg % jpeg_q))
    print(jpeg_msssim)
    cost = lambda q: avifenc(0, q)
    avif_q = search(1, 63, cost, jpeg_msssim)
    best_q = (avif_q, avif_q)
    best_sz = os.stat(test_avif % best_q).st_size
    for off in range(1, 8):
        cost = lambda q: avifenc(off, q)
        avif_q = search(max(1, avif_q - off), min(63 - off, avif_q + off),
                        cost, jpeg_msssim)
        sz = os.stat(test_avif % (avif_q, avif_q + off)).st_size
        if sz < best_sz:
            best_sz, best_q = sz, (avif_q, avif_q + off)
    return (test_jpg % jpeg_q), (test_avif % best_q)


def main(test_y4m):
    basename = os.path.basename(test_y4m)
    print(
        match_image(test_y4m,
                    test_y4m=basename + '.png.y4m',
                    test_png=basename + '.png',
                    test_jpg=basename + '-%s.jpg',
                    test_avif=basename + '-%s-%s.avif'))


main(sys.argv[1])