Patch nclx data into AVIF file header for specified Primaries / TRC / YUV Matrix

patch_cicp.py

import binascii
import os
import re
from logging import getLogger

logger = getLogger(__name__)

def patch_avif_header_cicp(input_file, P=9, T=16, M=9, output_file=None, preset=None, dryrun=False):
    """Patch nclx data into AVIF file header for specified Primaries / TRC / YUV Matrix Coeff enums."""

    CICP_ENCODINGS = {
        # primaries, trc, matrix
        'bt2100_pq':   (9, 16, 9),
        'bt2100_hlg':  (9, 18, 9),
        'bt709_srgb':  (1, 13, 1),
        'bt709_g22':   (1, 4, 1),
        'bt709_bt709': (1, 1, 1),
        'p3d65_pq':    (12, 16, 1),
        'p3d65_srgb':  (12, 13, 1),
        'p3d65_g22':   (12, 4, 1),
    }

    def cicp_to_bytes(P=1, T=13, M=1):
        return binascii.unhexlify("%04x%04x%04x" % (P, T, M))

    def bytes_to_cicp(v):
        v = binascii.hexlify(v)
        return [int(i, 16) for i in [v[0:4], v[5:8], v[9:12]]]

    output_file = output_file or input_file
    in_place = (input_file == output_file)
    P, T, M = CICP_ENCODINGS.get(preset, [P, T, M])
    ptrn, offs, cutlen = b"colrnclx", len(b"colrnclx"), 2 * len((P, T, M))

    with open(input_file, 'r+b' if in_place else 'rb') as fd_in:
        data = fd_in.read()
        match = re.search(ptrn, data)

        logger.debug('Patch %s: CICP: %s --> %s' % (
                os.path.basename(input_file),
                tuple(bytes_to_cicp(data[match.start() + offs:match.start() + offs + cutlen])),
                (P, T, M)))

        if not dryrun:
            if in_place:
                fd_in.seek(match.start()+offs)
                fd_in.write(cicp_to_bytes(P,T,M))
            else:
                with open(output_file, 'wb') as fd_out:
                    patched = data[:match.start() + offs] + cicp_to_bytes(P, T, M) + data[match.start() + offs + cutlen:]
                    fd_out.write(patched)

    logger.info(f'Patched CICP({P},{T},{M}): {output_file or input_file}')

    return output_file or input_file