Converts a YUV image to a PNG image using FFMPEG and Multithreading

yuv_to_png.py

import argparse
import multiprocessing
import os
import re
import subprocess
from functools import partial

import tqdm as tqdm

import PixelFormats


def convert_yuv_to_png(image_path: str, size: str = None, in_pixel_format: str = None, out_pixel_format: str = "rgb24",
                       save_output_to_file: bool = True):
    """
    Converts a YUV image to a PNG image.
    @attention: The PNG format is lossless for RGB24 image data. However the conversion from YUV to RGB24 is not
    lossless, as the two formats quantize the color space differently.
    @see https://dsp.stackexchange.com/questions/41083/loss-in-converting-yuv-to-png
    @param image_path: Image path pointing to the image to convert.
    @param size: The original image size. You can pass it as a string with the format <width>x<height>. Defaults to
    size format contained in file name and surrounded by underscores, e.g. `<filename>_<width>x<height>_<filename>.yuv`.
    @param in_pixel_format: The input pixel format. Defaults to pixel format name contained in file name and
    surrounded by underscores, e.g. `<filename>_<yuv420p>_<filename>.yuv`.
    @param out_pixel_format: The output pixel format. Run `ffmpeg -pix_fmts` for available formats. Defaults to `rgb24`.
    @param save_output_to_file: If true then the output of the ffmpeg is piped to a file. Actually two files are
    created one with stdout messages and another one with stderr messages. Defaults to True.
    @return:
    """
    if not os.path.isfile(image_path):
        raise FileNotFoundError(f"File not found: {image_path}.")

    # If size tuple is not passed then search on file name for _<width>x<height>_
    if size is None:
        match = re.search(r"_([\d]+x[\d]+)_", os.path.basename(image_path))
        if match:
            width, height = match.group().split("x")
            width = int(width.replace("_", ""))
            height = int(height.replace("_", ""))
        else:
            raise NameError("Unable to parse width and height. Expected yuv file name to be <some random "
                            "name>_<width>x<height>_<pixel format>_<other info>.yuv, got "
                            f"{os.path.basename(image_path)}")
    else:
        width, height = size.split("x")

    # If pixel format is not passed then search on file name for _<pixel format>_
    if in_pixel_format is None:
        for name in PixelFormats.all_input_formats():
            if "_" + name + "_" in os.path.basename(image_path):
                in_pixel_format = name
        if in_pixel_format is None:
            raise NameError("Unable to parse pixel format. Expected yuv file name to be <some random "
                            "name>_<width>x<height>_<pixel format>_<other info>.yuv, got "
                            f"{os.path.basename(image_path)}")

    # Defines the output file name (Saved in the same path as the original image)
    out_image_path = os.path.join(
        os.path.dirname(image_path),
        os.path.splitext(os.path.basename(image_path))[0] + ".png"
    )

    # From http://www.libpng.org/pub/png/book/chapter08.html:
    # RGB (truecolor) PNGs, like grayscale with alpha, are supported in only two depths: 8 and 16 bits per sample,
    # corresponding to 24 and 48 bits per pixel.
    cmd = [
        "ffmpeg",
        "-y",  # Allows overwriting
        "-s", f"{width}x{height}",
        "-pix_fmt", in_pixel_format,
        "-i", image_path,  # Sets input
        "-f", "image2",  # From https://ffmpeg.org/ffmpeg-formats.html#image2-2: The image file muxer writes video
        # frames to image files.
        "-pix_fmt", out_pixel_format,  # Sets pixel format. See available pixel formats with the command `ffmpeg
        # -pix_fmts`
        out_image_path  # Sets output file name
    ]

    result = subprocess.run(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE
    )

    if save_output_to_file:
        save_path_no_ext = os.path.join(
            os.path.dirname(image_path),
            os.path.splitext(os.path.basename(image_path))[0] + "_yuv_to_png"
        )
        with open(f"{save_path_no_ext}.stdout", "w") as fwriter:
            fwriter.write(result.stdout.decode("utf-8"))
        with open(f"{save_path_no_ext}.stderr", "w") as fwriter:
            fwriter.write(result.stderr.decode("utf-8"))
    return out_image_path


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Converts PNG images to YUV 420.",
                                     epilog="e.g. python yuv_to_png.py ../datasets/cityscapes/ | e.g. "
                                            "python yuv_to_png.py ../datasets/cityscapes/ -s train -fn _bs_decoded.yuv "
                                            "-go")
    parser.add_argument("dataset", type=str, help="Path to dataset.")
    parser.add_argument("-j", "--jobs", type=int, default=multiprocessing.cpu_count(),
                        help=f"Number of parallel jobs. Defauls to the total number of CPUs.")
    parser.add_argument("-is", "--in_size", type=str, default=None,
                        help="Input images size. Passed as a string in the format <width>x<height>. Defaults to "
                             "size format contained in file name and surrounded by underscores, "
                             "e.g. `<filename>_<width>x<height>_<filename>.yuv`.")
    parser.add_argument("-ipf", "--in_pix_fmt", type=str, default=None,
                        help="The input pixel format. Defaults to pixel format name contained in file name and "
                             "surrounded by underscores, e.g. `<filename>_<yuv420p>_<filename>.yuv`")
    parser.add_argument("-opf", "--out_pix_fmt", type=str, default="rgb24",
                        help="The output pixel format. Run `ffmpeg -pix_fmts` for available formats. Defaults to "
                             "`rgb24`.")
    parser.add_argument("-s", "--subset", type=str, default="",
                        help="Only files under `subset` are selected. Defaults to None (all files are selected).")
    parser.add_argument("-fn", "--filter_name", type=str, default="",
                        help="Only files under `subset` are selected. Defaults to None (all files are selected).")
    parser.add_argument("-go", "--gen_output", action="store_true",
                        help="If passed output of ffmpeg is stored in files.")
    args = parser.parse_args()

    files = []
    for dirpath, dirnames, filenames in os.walk(args.dataset):  # noqa
        for file in filenames:
            if os.path.splitext(file)[-1] == ".yuv" and \
               args.filter_name in file and \
               args.subset in os.path.basename(os.path.normpath(dirpath)):
                files.append(
                    os.path.join(
                        dirpath,
                        file
                    )
                )

    if not len(files):
        raise FileNotFoundError(f"No .yuv files found under {args.dataset}.")

    pool = multiprocessing.Pool(processes=args.jobs)

    partial_func = partial(
        convert_yuv_to_png,
        size=args.in_size,
        in_pixel_format=args.in_pix_fmt,
        out_pixel_format=args.out_pix_fmt,
        save_output_to_file=args.gen_output
    )
    for _ in tqdm.tqdm(pool.imap(partial_func, files), total=len(files)):
        pass
    print(">> done")