Skip to content

Commit

Permalink
COM-6142 (#187)
Browse files Browse the repository at this point in the history
COM-6142: Add AV1 argon test suite and generator

- created new test suite generator script
- modified test suite structure to accomodate this special case
- added the generated test suite
- tested to work on Ubuntu 20, 22 and MacOS 14.4
  • Loading branch information
mcesariniflu authored Oct 17, 2024
1 parent 5d70678 commit 9ac8828
Show file tree
Hide file tree
Showing 5 changed files with 25,747 additions and 5 deletions.
2 changes: 1 addition & 1 deletion .pylintrc
Original file line number Diff line number Diff line change
Expand Up @@ -300,7 +300,7 @@ max-bool-expr=5
max-branches=12

# Maximum number of locals for function / method body.
max-locals=15
max-locals=18

# Maximum number of parents for a class (see R0901).
max-parents=7
Expand Down
8 changes: 6 additions & 2 deletions fluster/codec.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,17 +37,21 @@ class OutputFormat(Enum):

NONE = "None"
YUV420P = "yuv420p"
YUV422P = "yuv422p"
YUV420P10LE = "yuv420p10le"
YUV420P12LE = "yuv420p12le"
YUV422P = "yuv422p"
YUV422P10LE = "yuv422p10le"
YUV422P12LE = "yuv422p12le"
YUV444P = "yuv444p"
YUV444P10LE = "yuv444p10le"
YUV444P12LE = "yuv444p12le"
YUV444P16LE = "yuv444p16le"
GBRP = "gbrp"
GBRP10LE = "gbrp10le"
GBRP12LE = "gbrp12le"
GBRP14LE = "gbrp14le"
GRAY = "gray"
GRAY10LE = "gray10le"
GRAY12LE = "gray12le"
GRAY16LE = "gray16le"
GBRP14LE = "gbrp14le"
UNKNOWN = "Unknown"
80 changes: 78 additions & 2 deletions fluster/test_suite.py
Original file line number Diff line number Diff line change
Expand Up @@ -224,6 +224,51 @@ def _download_worker(ctx: DownloadWork) -> None:
if not ctx.keep_file:
os.remove(dest_path)

@staticmethod
def _download_worker_av1_argon(ctx: DownloadWork) -> None:
"""Download and extract a av1 argon test vector"""
test_vector = ctx.test_vector
dest_dir = os.path.join(ctx.out_dir, ctx.test_suite_name)
dest_path = os.path.join(dest_dir, os.path.basename(test_vector.source))
if not os.path.exists(dest_dir):
os.makedirs(dest_dir)
# Catch the exception that download may throw to make sure pickle can serialize it properly
# This avoids:
# Error sending result: '<multiprocessing.pool.ExceptionWithTraceback object at 0x7fd7811ecee0>'.
# Reason: 'TypeError("cannot pickle '_io.BufferedReader' object")'
if not os.path.exists(dest_path):
print(f"\tDownloading test vector {test_vector.name} from {dest_dir}")
for i in range(ctx.retries):
try:
exception_str = ""
utils.download(test_vector.source, dest_dir)
except urllib.error.URLError as ex:
exception_str = str(ex)
print(
f"\tUnable to download {test_vector.source} to {dest_dir}, {exception_str}, retry count={i+1}"
)
continue
except Exception as ex:
raise Exception(str(ex)) from ex
break

if exception_str:
raise Exception(exception_str)
if test_vector.source_checksum != "__skip__":
checksum = utils.file_checksum(dest_path)
if test_vector.source_checksum != checksum:
raise Exception(
f"Checksum error for test vector '{test_vector.name}': '{checksum}' instead of "
f"'{test_vector.source_checksum}'"
)
if utils.is_extractable(dest_path):
print(f"\tExtracting test vector {test_vector.name} to {dest_dir}")
utils.extract(
dest_path,
dest_dir,
file=test_vector.input_file if not ctx.extract_all else None,
)

def download(
self,
jobs: int,
Expand All @@ -236,8 +281,25 @@ def download(
"""Download the test suite"""
if not os.path.exists(out_dir):
os.makedirs(out_dir)
if self.name == "AV1_ARGON_VECTORS":
# Only one job to download the zip file for Argon.
jobs = 1
dest_dir = os.path.join(out_dir, self.name)
test_vector_key = self.test_vectors[list(self.test_vectors)[0]].source
dest_folder = os.path.splitext(os.path.basename(test_vector_key))[0]
dest_path = os.path.join(dest_dir, dest_folder)
if (
verify
and os.path.exists(dest_path)
and self.test_vectors[test_vector_key].source_checksum
== utils.file_checksum(dest_path)
):
# Remove file only in case the input file was extractable.
# Otherwise, we'd be removing the original file we want to work
# with every even time we execute the download subcommand.
if utils.is_extractable(dest_path) and not keep_file:
os.remove(dest_path)
print(f"Downloading test suite {self.name} using {jobs} parallel jobs")

with Pool(jobs) as pool:

def _callback_error(err: Any) -> None:
Expand All @@ -255,9 +317,13 @@ def _callback_error(err: Any) -> None:
test_vector,
retries,
)
if self.name == "AV1_ARGON_VECTORS":
download_worker = self._download_worker_av1_argon
else:
download_worker = self._download_worker
downloads.append(
pool.apply_async(
self._download_worker,
download_worker,
args=(dwork,),
error_callback=_callback_error,
)
Expand All @@ -269,6 +335,16 @@ def _callback_error(err: Any) -> None:
if not job.successful():
sys.exit("Some download failed")

if self.name == "AV1_ARGON_VECTORS":
if not dwork.keep_file:
os.remove(
os.path.join(
dwork.out_dir,
dwork.test_suite_name,
os.path.basename(dwork.test_vector.source),
)
)

print("All downloads finished")

@staticmethod
Expand Down
207 changes: 207 additions & 0 deletions scripts/gen_av1_argon.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,207 @@
#!/usr/bin/env python3

# Fluster - testing framework for decoders conformance
# Copyright (C) 2020-2024, Fluendo, S.A.
# Author: Martin Cesarini <[email protected]>, Fluendo, S.A.
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public License
# as published by the Free Software Foundation, either version 3
# of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library. If not, see <https://www.gnu.org/licenses/>.

import argparse
import os
import re
import subprocess
import sys
import urllib.error
import zipfile

# pylint: disable=wrong-import-position
sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
from fluster import utils
from fluster.codec import Codec, OutputFormat
from fluster.test_suite import TestSuite, TestVector

# pylint: enable=wrong-import-position

ARGON_URL = "https://storage.googleapis.com/downloads.aomedia.org/assets/zip/"


class AV1ArgonGenerator:
"""Generates a test suite from the conformance bitstreams"""

def __init__(
self,
name: str,
suite_name: str,
codec: Codec,
description: str,
site: str,
use_ffprobe: bool = False,
):
self.name = name
self.suite_name = suite_name
self.codec = codec
self.description = description
self.site = site
self.use_ffprobe = use_ffprobe

def generate(self, download):
"""Generates the test suite and saves it to a file"""
output_filepath = os.path.join(self.suite_name + ".json")
extract_folder = "resources"
test_suite = TestSuite(
output_filepath,
extract_folder,
self.suite_name,
self.codec,
self.description,
dict(),
)
os.makedirs(extract_folder, exist_ok=True)
# Download the zip file
source_url = self.site + self.name
if download:
print(f"Download test suite archive from {source_url}")
try:
utils.download(source_url, extract_folder)
except urllib.error.URLError as ex:
exception_str = str(ex)
print(
f"\tUnable to download {source_url} to {extract_folder}, {exception_str}"
)
except Exception as ex:
raise Exception(str(ex)) from ex

# Unzip the file
test_vector_files = []
with zipfile.ZipFile(extract_folder + "/" + self.name, "r") as zip_ref:
print(f"Unzip files from {self.name}")
for file_info in zip_ref.namelist():
# Extract test vector files
if file_info.endswith(".obu"):
zip_ref.extract(file_info, extract_folder)
test_vector_files.append(file_info)

# Extract md5 files
if (
file_info.endswith(".md5")
and "md5_ref/" in file_info
and "layers/" not in file_info
):
zip_ref.extract(file_info, extract_folder)

# Create the test vector and test suite
print("Creating test vectors and test suite")
source_checksum = utils.file_checksum(extract_folder + "/" + self.name)
for file in test_vector_files:
filename = os.path.splitext(os.path.basename(file))[0]
# ffprobe execution
if self.use_ffprobe:
full_path = os.path.abspath(extract_folder + "/" + file)
ffprobe = utils.normalize_binary_cmd("ffprobe")
command = [
ffprobe,
"-v",
"error",
"-select_streams",
"v:0",
"-show_entries",
"stream=pix_fmt",
"-of",
"default=nokey=1:noprint_wrappers=1",
full_path,
]
try:
result = utils.run_command_with_output(command).splitlines()
pix_fmt = result[0]
if pix_fmt == "unknown":
pix_fmt = "Unknown"
except subprocess.CalledProcessError:
pix_fmt = "None"

# Processing md5 files
md5_file_to_find = os.path.splitext(filename)[0] + ".md5"
full_path_split = full_path.split("/")
md5_directory_path = (
"/".join(full_path_split[: len(full_path_split) - 2])
+ "/"
+ "md5_ref"
)
md5_file_path = os.path.join(md5_directory_path, md5_file_to_find)

# Check the .md5 file and get checksum
if os.path.exists(md5_file_path):
try:
result_checksum = self._fill_checksum_argon(md5_file_path)
except Exception as ex:
print("MD5 does not match")
raise ex
else:
try:
result_checksum = utils.file_checksum(full_path)
except Exception as ex:
print("MD5 cannot be calculated")
raise ex

# Add data to the test vector and the test suite
test_vector = TestVector(
filename,
source_url,
source_checksum,
file,
OutputFormat[pix_fmt.upper()],
result_checksum,
)
test_suite.test_vectors[filename] = test_vector

test_suite.to_json_file(output_filepath)
print("Generate new test suite: " + test_suite.name + ".json")

@staticmethod
def _fill_checksum_argon(dest_dir):
checksum_file = dest_dir
if checksum_file is None:
raise Exception("MD5 not found")
with open(checksum_file, "r") as checksum_file:
regex = re.compile(r"([a-fA-F0-9]{32,}).*\.(yuv|rgb|gbr)")
lines = checksum_file.readlines()
if any((match := regex.match(line)) for line in lines):
result = match.group(1)[:32].lower()
else:
result = -1
# Assert that we have extracted a valid MD5 from the file
assert (
len(result) == 32
and re.search(r"^[a-fA-F0-9]{32}$", result) is not None
), f"{result} is not a valid MD5 hash"
return result


if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"--skip-download",
help="skip extracting tarball",
action="store_true",
default=False,
)
args = parser.parse_args()
generator = AV1ArgonGenerator(
"argon_coveragetool_av1_base_and_extended_profiles_v2.1.1.zip",
"AV1_ARGON_VECTORS",
Codec.AV1,
"AV1 Argon Streams",
ARGON_URL,
True,
)
generator.generate(not args.skip_download)
Loading

0 comments on commit 9ac8828

Please sign in to comment.