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

png_to_yuv.py

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

import tqdm
from cv2 import cv2 as cv


def get_png_size(image_path):
    return cv.imread(image_path).shape[:2]


def convert_png_to_yuv(image_path: str, pixel_format: str = "yuv420p", save_output_to_file: bool = True):
    """
    Convert a PNG image to a YUV image.
    @param image_path: The PNG image path. Defaults to yuv420p.
    @param pixel_format: The pixel format to use in conversion. See ffmpeg -pix_fmts or PixelFormats class.
    @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: None
    """
    if not os.path.isfile(image_path):
        raise FileNotFoundError(f"File not found: {image_path}.")

    height, width = get_png_size(image_path)
    # pixel_format_int = pixel_format.replace("yuv", "").replace("p", "")
    out_image_path = os.path.join(
        os.path.dirname(image_path),
        os.path.splitext(os.path.basename(image_path))[0] + f"_{width}x{height}_{pixel_format}.yuv"
    )
    cmd = [
        "ffmpeg",
        "-y",  # Allows overwritting
        "-i", image_path,  # Sets input
        "-pix_fmt", 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] + "_ffmpeg"
        )
        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 png_to_yuv.py ../datasets/cityscapes/ | e.g. "
                                            "python png_to_yuv.py ../datasets/cityscapes/ -pf yuv444p -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("-pf", "--pix_fmt", type=str, default="yuv420p",
                        help="Pixel format. Run `ffmpeg -pix_fmts` for available formats. Defaults to yuv420p.")
    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] == ".png":
                files.append(
                    os.path.join(
                        dirpath,
                        file
                    )
                )

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

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

    # pool.map is blocking, pool.map_async is non blocking!
    # pool.imap iterates trough the argsuments instead of copiyng a chunck of it.
    # pool.imap_unordered iterates trough the argsuments instead of copiyng a chunck of it but the results are not
    # retrived by the order of the passed args.
    partial_func = partial(convert_png_to_yuv, pixel_format=args.pix_fmt, save_output_to_file=args.gen_output)
    for _ in tqdm.tqdm(pool.imap(partial_func, files), total=len(files)):
        pass
    print(">> done")