add script to split complexity csv files by levels

#! /usr/bin/env python3

# (C) 2022-2025 IVAS codec Public Collaboration with portions copyright Dolby International AB, Ericsson AB,

# Fraunhofer-Gesellschaft zur Foerderung der angewandten Forschung e.V., Huawei Technologies Co. LTD.,

# Koninklijke Philips N.V., Nippon Telegraph and Telephone Corporation, Nokia Technologies Oy, Orange,

# Panasonic Holdings Corporation, Qualcomm Technologies, Inc., VoiceAge Corporation, and other

# contributors to this repository. All Rights Reserved.

# This software is protected by copyright law and by international treaties.

# The IVAS codec Public Collaboration consisting of Dolby International AB, Ericsson AB,

# Fraunhofer-Gesellschaft zur Foerderung der angewandten Forschung e.V., Huawei Technologies Co. LTD.,

# Koninklijke Philips N.V., Nippon Telegraph and Telephone Corporation, Nokia Technologies Oy, Orange,

# Panasonic Holdings Corporation, Qualcomm Technologies, Inc., VoiceAge Corporation, and other

# contributors to this repository retain full ownership rights in their respective contributions in

# the software. This notice grants no license of any kind, including but not limited to patent

# license, nor is any license granted by implication, estoppel or otherwise.

# Contributors are required to enter into the IVAS codec Public Collaboration agreement before making

# contributions.

# This software is provided "AS IS", without any express or implied warranties. The software is in the

# development stage. It is intended exclusively for experts who have experience with such software and

# solely for the purpose of inspection. All implied warranties of non-infringement, merchantability

# and fitness for a particular purpose are hereby disclaimed and excluded.

# Any dispute, controversy or claim arising under or in relation to providing this software shall be

# submitted to and settled by the final, binding jurisdiction of the courts of Munich, Germany in

# accordance with the laws of the Federal Republic of Germany excluding its conflict of law rules and

# the United Nations Convention on Contracts on the International Sales of Goods.

import pandas as pd

import numpy as np

import argparse

import pathlib

LEVELS_2_MAX_BR = {

    "level_1": (0, 80),

    "level_2": (96, 192),

    "level_3": (256, 512),

}

def main(csv_files: list):

    for csv_file in csv_files:

        df = pd.read_csv(csv_file, delimiter=";")

        mask_vbr = df["conf"].str.contains("RS")

        df_vbr = df[mask_vbr]

        filename_vbr = f"{csv_file.stem}_rate_sw{csv_file.suffix}"

        df_vbr.to_csv(

            filename_vbr,

            index=False,

            sep=";",

        )

        df_cbr = df[np.logical_not(mask_vbr)]

        df_cbr["bitrate"] = df_cbr["conf"].str.extract(r"@(.*) kbps").astype("float")

        for lvl_suffix, (min_br, max_br) in LEVELS_2_MAX_BR.items():

            mask_lvl = np.logical_and(

                df_cbr["bitrate"] >= min_br, df_cbr["bitrate"] <= max_br

            )

            df_lvl = df_cbr[mask_lvl]

            filename_lvl = f"{csv_file.stem}_{lvl_suffix}{csv_file.suffix}"

            df_lvl.to_csv(

                filename_lvl,

                index=False,

                sep=";",

            )

if __name__ == "__main__":

    parser = argparse.ArgumentParser()

    parser.add_argument(

        "csv_files",

        nargs="+",

        help="CSV files to split by complexity levels.",

        type=pathlib.Path,

    )

    args = parser.parse_args()

    main(args.csv_files)