Python based CLI for Encoding and Decoding a whole Dataset using HM and Multihreading

hm_encode_decode.py

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

import tqdm


def encode(hm_encoder_path: str, image_path: str, config_path: str, qp: int, output_dir: str = None,
           verbosity: int = 0, gen_recon: bool = True, gen_bitstream: bool = True,
           gen_report: bool = False, gen_reportpic: bool = False, save_output_to_file: bool = True, **kwargs):
    """
    Encodes a YUV image using HEVC encoder software. Expecting the image name to be: <some random
    name>_<width>x<height>_<pixel format>.yuv
    @param hm_encoder_path: The HM encoder path, usually is TAppEncoderStatic under bin directory.
    @param image_path: The yuv image to encode.
    @param config_path: The config path.
    @param qp: QP to use to encode.
    @param output_dir: All files generated by this script will be saved under <output_dir>. Defauls to <image_path> dir.
    @param verbosity: The verbosity level applied to all generated files.
    @param gen_recon: If True recon file is saved with `_rec.yuv` appended to the input path.
    @param gen_bitstream: If True recon file is saved with `_bs.bin` appended to the input path.
    @param gen_report: If True a summary/report is generated and saved with `_report.txt` appended to the input path.
    @param gen_reportpic: If True a summary/report per picture is generated and saved with `_reportpic<PICTURE
    TYPE>.txt` appended to the input path.
    @param save_output_to_file: If True the output generated by the HM encoder tool is saved.
    @param kwargs: Additional parameters that are added to the default command. The additonal options are appended to
    the default command: `<hm_encoder_path> -i <image_path> -c <config_path> -q <qp> -o <out_recon_path> -b
    <out_bitstream_path> -wdt <width> -hgt <height> -fr 1 -f 1 -v <verbosity> --InputChromaFormat=<pixel_format>
    --InputBitDepth=8 --ConformanceWindowMode=1 [--SummaryOutFilename=<out_summary_path>]
    [--SummaryPicFilenameBase=<out_summarypic_path>] [--SummaryVerboseness=<verbosity>]`.
    @return: The recon file path and the bitstream file path. (out_recon_path, out_bitstream_path)
    """
    match = re.search(r"([\d]+x[\d]+)", os.path.basename(image_path))
    if match:
        width, height = match.group().split("x")
    else:
        return None

    if not output_dir:
        output_dir = os.path.dirname(image_path)

    os.makedirs(output_dir, exist_ok=True)

    pixel_format = os.path.splitext(os.path.basename(image_path))[0].split("_")[-1].replace("yuv", "").replace("p", "")
    out_recon_path = os.path.join(
        output_dir,
        os.path.splitext(os.path.basename(image_path))[0] + "_rec.yuv"
    ) if gen_recon else "/dev/null"
    out_bitstream_path = os.path.join(
        output_dir,
        os.path.splitext(os.path.basename(image_path))[0] + "_bs.bin"
    ) if gen_bitstream else "/dev/null"
    out_summary_path = os.path.join(
        output_dir,
        os.path.splitext(os.path.basename(image_path))[0] + "_report.txt"
    )
    out_summarypic_path = os.path.join(
        output_dir,
        os.path.splitext(os.path.basename(image_path))[0] + "_reportpic"
    )

    cmd = [
        hm_encoder_path,  # High Efficiency Video Coding (HEVC): HM Encoder
        "-i", image_path,  # Original YUV input file name
        "-c", config_path,  # Configuration file name
        "-q", str(qp),  # Qp value
        "-o", out_recon_path,  # Reconstructed YUV output file name
        "-b", out_bitstream_path,  # Bitstream output file name
        "-wdt", width,  # Source picture width
        "-hgt", height,  # Source picture height
        "-fr", "1",  # Frame rate
        "-f", "1",  # Number of frames to be encoded (default=all)
        f"--InputChromaFormat={pixel_format}",  # InputChromaFormatIDC
        "--InputBitDepth=8",  # Bit-depth of input file
        "--ConformanceWindowMode=1"  # Window conformance mode (0: no window, 1:automatic padding (default),
                                     # 2:padding parameters specified, 3:conformance window parameters specified
    ]
    if gen_report:
        cmd += [
            f"--SummaryOutFilename={out_summary_path}"  # Filename to use for producing summary output file. If
            # empty, do not produce a file.
        ]
    if gen_reportpic:
        cmd += [
            f"--SummaryPicFilenameBase={out_summarypic_path}"  # Base filename to use for producing summary picture
            # output files. The actual filenames used will have
            # I.txt, P.txt and B.txt appended. If empty,
            # do not produce a file.
        ]
    if gen_report or gen_reportpic:
        cmd += [
            f"--SummaryVerboseness={verbosity}",  # Specifies the level of the verboseness of the text output.
        ]

    # Add extra parameters to HM TAppDecoderStatic command
    for k, v in kwargs.items():  # noqa
        cmd += [
            str(k), str(v)
        ]

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

    if save_output_to_file:
        save_path_no_ext = os.path.join(
            output_dir,
            os.path.splitext(os.path.basename(image_path))[0] + "_hmencoder"
        )
        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_recon_path, out_bitstream_path


def decode(hm_decoder_path: str, bitstream_path: str, output_dir: str = None,
           gen_decoded: bool = True, save_output_to_file: bool = True, **kwargs):
    """
    Decodes a bitstream file using HM.
    @param hm_decoder_path: HM decoder path, usually is TAppDecoderStatic in bin directory.
    @param bitstream_path: The bitstream path.
    @param output_dir: All files generated by this script will be saved under <output_dir>. Defauls to <image_path> dir.
    @param gen_decoded: If True the decoded file will be saved with `_decoded.yuv` appended to the bitstream name.
    @param save_output_to_file: If True the output generated by the HM decoder tool is saved.
    @param kwargs: Additional parameters that are added to the default command. The additonal options are appended to
    the default command: `<hm_decoder_path> -b <bitstream_path> -o <out_recon_path> -d 8`.
    next to the
    @return: The decoded/reconstructed yuv file name. (out_recon_path)
    @warning bit_depth: The bitdepth of the decoded/reconstructed bitstream. This tool always uses 8 bits
    to be compliant with the encode command.
    """

    if not output_dir:
        output_dir = os.path.dirname(bitstream_path)

    os.makedirs(output_dir, exist_ok=True)

    out_recon_path = os.path.join(
        output_dir,
        os.path.splitext(os.path.basename(bitstream_path))[0] + "_decoded.yuv"
    ) if gen_decoded else "/dev/null"

    cmd = [
        hm_decoder_path,
        "-b", bitstream_path,
        "-o", out_recon_path,
        "-d", "8"
    ]

    # Add extra parameters to HM TAppDecoderStatic command
    for k, v in kwargs.items():  # noqa
        cmd += [
            str(k), str(v)
        ]

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

    if save_output_to_file:
        save_path_no_ext = os.path.join(
            output_dir,
            os.path.splitext(os.path.basename(bitstream_path))[0] + "_hmdecoder"
        )
        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_recon_path


def encode_decode(
        hm_encoder_path: str, image_path: str, config_path: str, qp: int,      # encode args
        hm_decoder_path: str,                                                  # decode args
        verbosity: int = 0, gen_recon: bool = True, gen_report: bool = False,   # encode optional args
        gen_reportpic: bool = False,                                            # encode optional args
        gen_decoded: bool = True,                                               # decode optional args
        output_dir: str = None, save_output_to_file: bool = True                # common optional args
):
    """
    Encodes and Decodes an YUV image. See encode and decode function docs for more details.
    """
    recon_path, bitstream_path = encode(
        hm_encoder_path=hm_encoder_path,
        image_path=image_path,
        config_path=config_path,
        qp=qp,
        output_dir=output_dir,
        verbosity=verbosity,
        gen_recon=gen_recon,
        gen_bitstream=True,
        gen_report=gen_report,
        gen_reportpic=gen_reportpic,
        save_output_to_file=save_output_to_file
    )
    decoded_path = decode(
        hm_decoder_path=hm_decoder_path,
        bitstream_path=bitstream_path,
        output_dir=output_dir,
        gen_decoded=gen_decoded,
        save_output_to_file=save_output_to_file
    )
    return recon_path, bitstream_path, decoded_path


def _encode_decode_wrapper(_image_path, _args):
    """
    Needed to allow multiprocessing since the <image_path> parameter will be the one passed with multiple values.
    """
    return encode_decode(
        hm_encoder_path=_args.hm_encoder,
        image_path=_image_path,
        config_path=_args.config_path,
        qp=_args.qp,
        hm_decoder_path=_args.hm_decoder,
        verbosity=_args.verbosity,
        gen_recon=_args.gen_recon,
        gen_report=_args.gen_summary,
        gen_reportpic=_args.gen_summarypic,
        gen_decoded=_args.no_gen_decoded,
        output_dir=_args.output_dir,
        save_output_to_file=_args.gen_output
    )


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Encodes YUV images using HM in a multithread scheme.",
                                     epilog="e.g. python hm_encode_decode.py "
                                            "../repos/HM/bin/TAppEncoderStatic "
                                            "../datasets/cityscapes/ "
                                            "../repos/HM/cfg/encoder_intra_main.cfg "
                                            "47 "
                                            "../repos/HM/bin/TAppDecoderStatic")
    parser.add_argument("hm_encoder", type=str, help="Path to HM Encoder (usually refered as TAppEncoderStatic).")
    parser.add_argument("dataset", type=str, help="Path to dataset subset.")
    parser.add_argument("config_path", type=str, help="Path to HM config.")
    parser.add_argument("qp", type=int, help="QP to use when encoding pictures.")
    parser.add_argument("hm_decoder", type=str, help="Path to HM Decoder (usually refered as TAppDecoderStatic).")
    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("-o", "--output_dir", type=str, default=None,
                        help=f"If passed then its used to save all output files, else the same directory of the "
                             f"images is used.")
    parser.add_argument("-v", "--verbosity", type=int, default=0, help=f"Verbosity level.")
    parser.add_argument("-gr", "--gen_recon", action="store_true",
                        help="If passed recon files are saved.")
    parser.add_argument("-ngd", "--no_gen_decoded", action="store_false",
                        help="If passed decoded YUV from bitstream is NOT stored in files.")
    parser.add_argument("-gs", "--gen_summary", action="store_true",
                        help="If passed report/summary files are saved.")
    parser.add_argument("-gsp", "--gen_summarypic", action="store_true",
                        help="If passed report/summary picture files are saved.")
    parser.add_argument("-go", "--gen_output", action="store_true",
                        help="If passed output of HM 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":
                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)

    for _ in tqdm.tqdm(
            pool.imap(
                partial(_encode_decode_wrapper, _args=args), files
            ), total=len(files)
    ):
        pass

    print(">> done")