convert.py

#!/usr/bin/env python3
"""
Convert any YUV to any range or bitdepth
"""
import argparse

DATA_TABLE = {
    ("tv", 8): [[16, 128, 128], [219, 224, 224]],
    ("tv", 10): [[64, 512, 512], [876, 896, 896]],
    ("tv", 16): [[4096, 32768, 32768], [56064, 57344, 57344]],
    ("pc", 8): [[0, 128, 128], [255] * 3],
    ("pc", 10): [[0, 512, 512], [1023] * 3],
    ("pc", 16): [[0, 32768, 32768], [65535] * 3],
}


def convert(yuv, bitrange="tv", bitdepth=8, target=("pc", 10)):
    """
    Convert yuv to specific range and bitdepth
    """
    return [round(i) for i in to_integer(to_float(yuv, bitrange, bitdepth), *target)]


def to_float(yuv, bitrange="tv", bitdepth=8):
    """
    Convert integer yuv to float
    """
    return [(x - y) / z for x, y, z in zip(yuv, *DATA_TABLE[bitrange, bitdepth])]


def to_integer(yuv, bitrange="tv", bitdepth=8):
    """
    Convert float yuv to integer
    """
    return [x * z + y for x, y, z in zip(yuv, *DATA_TABLE[bitrange, bitdepth])]


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Convert integer YUV to any range and any bitdepth"
    )

    parser.add_argument("yuv", type=int, nargs=3)
    parser.add_argument("--bitdepth", type=int, default=8, help="original bitdepth")
    parser.add_argument(
        "--bitrange", type=str, default="tv", help="original range: tv or pc"
    )
    parser.add_argument("--target_bit", type=int, default=10, help="target bitdepth")
    parser.add_argument(
        "--target_range", type=str, default="pc", help="target range: tv or pc"
    )
    args = parser.parse_args()
    print(
        convert(
            args.yuv, args.bitrange, args.bitdepth, (args.target_range, args.target_bit)
        )
    )