Last active
December 2, 2021 17:54
-
-
Save pinxau1000/529c52c94e77560f93118becae536583 to your computer and use it in GitHub Desktop.
Python based CLI for Encoding and Decoding a whole Dataset using VTM and Multihreading
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import argparse | |
import multiprocessing | |
import os | |
import re | |
import subprocess | |
from functools import partial | |
import tqdm | |
def encode(vtm_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 VTM encoder software. Expecting the image name to be: <some random | |
name>_<width>x<height>_<pixel format>.yuv | |
@param vtm_encoder_path: The VTM encoder path, usually is EncoderAppStatic 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 VTM encoder tool is saved. | |
@param kwargs: Additional parameters that are added to the default command. The additonal options are appended to | |
the default command: `<vtm_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 = [ | |
vtm_encoder_path, # VVCSoftware: VTM 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) | |
"-v", str(verbosity), # Specifies the level of the verboseness | |
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 VTM DecoderAppStatic 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] + "_vtmencoder" | |
) | |
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(vtm_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 VTM. | |
@param vtm_decoder_path: VTM deocder path, usually is DecoderAppStatic 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 VTM decoder tool is saved. | |
@param kwargs: Additional parameters that are added to the default command. The additonal options are appended to | |
the default command: `<vtm_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 = [ | |
vtm_decoder_path, | |
"-b", bitstream_path, | |
"-o", out_recon_path, | |
"-d", "8" | |
] | |
# Add extra parameters to VTM DecoderAppStatic 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] + "_vtmdecoder" | |
) | |
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( | |
vtm_encoder_path: str, image_path: str, config_path: str, qp: int, # encode args | |
vtm_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( | |
vtm_encoder_path=vtm_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( | |
vtm_decoder_path=vtm_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( | |
vtm_encoder_path=_args.vtm_encoder, | |
image_path=_image_path, | |
config_path=_args.config_path, | |
qp=_args.qp, | |
vtm_decoder_path=_args.vtm_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 VTM in a multithread scheme.", | |
epilog="e.g. python vtm_encode_decode.py " | |
"../repos/VVCSoftware_VTM/bin/EncoderAppStatic " | |
"../datasets/cityscapes/ " | |
"../repos/VVCSoftware_VTM/cfg/encoder_intra_vtm.cfg " | |
"47 " | |
"../repos/VVCSoftware_VTM/bin/DecoderAppStatic") | |
parser.add_argument("vtm_encoder", type=str, help="Path to VTM Encoder (usually refered as EncoderAppStatic).") | |
parser.add_argument("dataset", type=str, help="Path to dataset subset.") | |
parser.add_argument("config_path", type=str, help="Path to VTM config.") | |
parser.add_argument("qp", type=int, help="QP to use when encoding pictures.") | |
parser.add_argument("vtm_decoder", type=str, help="Path to VTM Decoder (usually refered as DecoderAppStatic).") | |
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 VTM 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") |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment