COM-6142

.pylintrc

@@ -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

fluster/codec.py

@@ -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"

fluster/test_suite.py

@@ -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,
@@ -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:
@@ -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,
                     )
@@ -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

scripts/gen_av1_argon.py

@@ -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)