use formats from config file and add option to run only specific formats

#!/usr/bin/env python3

import argparse

import pandas as pd

import logging

import os

import re

import sys

import time

import json

from pathlib import Path

from concurrent.futures import ThreadPoolExecutor, as_completed

sys.path.append("./ivas-processing-scripts")

from ivas_processing_scripts.audiotools.wrappers.bs1770 import get_loudness

from ivas_processing_scripts.audiotools.audio import fromfile

from ivas_processing_scripts.utils import progressbar_update, spinner

from ivas_processing_scripts.utils import progressbar_update

logging.basicConfig(level=logging.WARNING)

if os.environ.get("CI") or not sys.stdout.isatty():

    sys.stdout.reconfigure(line_buffering=True)

INPUT_FOLDER_BASE = Path(__file__).parent.joinpath("testv", "pinknoise")

OUTPUT_FOLDER = Path(__file__).parent.parent.joinpath("out/dec")

FORMATS = [

    "MONO",

    "STEREO",

    "5_1",

    "5_1_2",

    "5_1_4",

    "7_1",

    "7_1_4",

    "ISM1",

    "ISM2",

    "ISM3",

    "ISM4",

    "MASA1DIR1",

    "MASA1DIR2",

    "MASA2DIR1",

    "MASA2DIR2",

    "FOA",

    "HOA2",

    "HOA3",

    # "ISM1SBA1",

    # "ISM2SBA1",

    # "ISM3SBA1",

    # "ISM4SBA1",

    # "ISM1SBA2",

    # "ISM2SBA2",

    # "ISM3SBA2",

    # "ISM4SBA2",

    # "ISM1SBA3",

    # "ISM2SBA3",

    # "ISM3SBA3",

    "ISM4SBA3",

    # "ISM1MASA1DIR1",

    # "ISM2MASA1DIR1",

    # "ISM3MASA1DIR1",

    # "ISM4MASA1DIR1",

    # "ISM1MASA1DIR2",

    # "ISM2MASA1DIR2",

    # "ISM3MASA1DIR2",

    # "ISM4MASA1DIR2",

    # "ISM1MASA2DIR1",

    # "ISM2MASA2DIR1",

    # "ISM3MASA2DIR1",

    # "ISM4MASA2DIR1",

    # "ISM1MASA2DIR2",

    # "ISM2MASA2DIR2",

    # "ISM3MASA2DIR2",

    "ISM4MASA2DIR2",

]

FORMAT_2_FILE = {

    fmt: path

    for fmt, path in zip(

        FORMATS, [INPUT_FOLDER_BASE.joinpath(f"{f}.wav") for f in FORMATS]

    )

CFG_TO_PROC_FMT = {

    "MONO": "MONO",

    "STEREO": "STEREO",

    "FOA": "FOA",

    "HOA2": "HOA2",

    "HOA3": "HOA3",

    "SBA": "HOA3",

    "MASA1TC": "MASA1DIR1",

    "MASA2TC": "MASA2DIR2",

    "5_1": "5_1",

    "5_1_2": "5_1_2",

    "5_1_4": "5_1_4",

    "7_1": "7_1",

    "7_1_4": "7_1_4",

    "ISM1": "ISM1",

    "ISM2": "ISM2",

    "ISM3": "ISM3",

    "ISM4": "ISM4",

    "OMASA_ISM1_1TC": "ISM1MASA1DIR1",

    "OMASA_ISM1_2TC": "ISM1MASA2DIR2",

    "OMASA_ISM2_1TC": "ISM2MASA1DIR1",

    "OMASA_ISM2_2TC": "ISM2MASA2DIR2",

    "OMASA_ISM3_1TC": "ISM3MASA1DIR1",

    "OMASA_ISM3_2TC": "ISM3MASA2DIR2",

    "OMASA_ISM4_1TC": "ISM4MASA1DIR1",

    "OMASA_ISM4_2TC": "ISM4MASA2DIR2",

    "OSBA_ISM1_FOA": "ISM1SBA1",

    "OSBA_ISM1_HOA2": "ISM1SBA2",

    "OSBA_ISM1_HOA3": "ISM1SBA3",

    "OSBA_ISM2_FOA": "ISM2SBA1",

    "OSBA_ISM2_HOA2": "ISM2SBA2",

    "OSBA_ISM2_HOA3": "ISM2SBA3",

    "OSBA_ISM3_FOA": "ISM3SBA1",

    "OSBA_ISM3_HOA2": "ISM3SBA2",

    "OSBA_ISM3_HOA3": "ISM3SBA3",

    "OSBA_ISM4_FOA": "ISM4SBA1",

    "OSBA_ISM4_HOA2": "ISM4SBA2",

    "OSBA_ISM4_HOA3": "ISM4SBA3",

}

PATTERN_BITRATE = re.compile(r"b([\d_]*|all)(_dtx)?_(swb|wb|fb)")

RESULT_OUTPUT_FILE = Path(__file__).parent.parent.joinpath("loudness.csv")

            return None

        output_audio = fromfile(outformat.upper(), outfile)

        output_loudness, _, output_loudness_format = get_loudness(output_audio)

        output_loudness, scale_factor, output_loudness_format = get_loudness(

            output_audio

        )

        return {

            "infile": infile.name,

            "dtx": dtx,

            "output_loudness": output_loudness,

            "output_loudness_format": output_loudness_format,

            "scale_factor": scale_factor,

        }

    except (AssertionError, ValueError) as e:

        print(f"\n⚠️  Skipping {outfile.stem}: {e}", file=sys.stderr)

        return None

def main():

    formats = FORMAT_2_FILE.keys()

def main(args):

    with open(args.config_json, "r") as f:

        config = json.load(f)

        config = config["inpaths"]

    FORMAT_2_FILE = {CFG_TO_PROC_FMT[k]: Path(v) for k, v in config.items()}

    valid_formats = [f for f in args.formats if f in CFG_TO_PROC_FMT]

    invalid_formats = [f for f in args.formats if f not in CFG_TO_PROC_FMT]

    if invalid_formats:

        print(f"The specified format(s) are invalid: {invalid_formats}")

        print(

            f"These do not match those in the provided config file: {args.config_json}"

        )

    if not valid_formats:

        print("No valid formats to process, exiting...")

        exit(-1)

    # map to proc scripts format names internally

    valid_formats = [CFG_TO_PROC_FMT[f] for f in valid_formats]

    results = []

    input_audio_cache = {}

    tasks = []

    for format in formats:

    for format in valid_formats:

        infile = FORMAT_2_FILE[format]

        output_folder = OUTPUT_FOLDER

                (outfile, infile, format, input_loudness, input_loudness_format)

            )

    print(f"Found {len(tasks)} files to process across {len(formats)} format(s)\n")

    print(

        f"Found {len(tasks)} files to process across {len(valid_formats)} format(s)\n"

    )

    if not len(tasks):

        print("Nothing to do, exiting...")

        exit(-1)

    # Process in parallel with progress bar

    print(f"⏳ Processing {len(tasks)} audio files:")

    start_time = time.time()

    completed = 0

    total = len(tasks)

            executor.submit(process_output_file, *task): task[0] for task in tasks

        }

        # Process with animated spinner and progress bar

        while futures:

            done = set()

        # Process with progress bar

        for future in as_completed(futures):

                done.add(future)

            outfile = futures[future]

            try:

            completed += 1

            progressbar_update(completed, total, width=50)

            # Remove completed futures

            for future in done:

                del futures[future]

            # Animate spinner even when waiting

            if futures:

                spinner()

    print()  # New line after progress bar

    # Save results

    elapsed = time.time() - start_time

    print(f"💾 Saving {len(results)} results to {RESULT_OUTPUT_FILE}")

    print(f"\n💾 Saving {len(results)} results to {RESULT_OUTPUT_FILE}")

    df = pd.DataFrame(results)

    df.to_csv(RESULT_OUTPUT_FILE, index=False)

    print(

        f"✅ Processed {len(results)}/{len(tasks)} files successfully in {elapsed:.1f}s"

    )

    print(f"   Average: {len(results) / elapsed:.1f} files/second")

    print(f"✅ Processed {len(results)}/{len(tasks)} files successfully")

if __name__ == "__main__":

    main()

    parser = argparse.ArgumentParser()

    parser.add_argument(

        "config_json",

        type=Path,

        default=Path(__file__).parent.joinpath("config", "ci_loudness_linux.json"),

    )

    parser.add_argument(

        "--formats",

        nargs="*",

        default=list(CFG_TO_PROC_FMT.keys()),

        help="Formats to process (default = %(default)s)",

    )

    args = parser.parse_args()

    main(args)

    y_ticks = np.arange(y_min, y_max + 5, 5)

    # get unique bitrates and sort them

    bitrates = sorted(df["bitrate"].unique())

    bitrates = sorted(filtered_df["bitrate"].unique())

    bitrate_to_idx = {br: idx for idx, br in enumerate(bitrates)}

    # bandwidth values

    bandwidths = df["bandwidth"].unique()

    bandwidths = filtered_df["bandwidth"].unique()

    # check for DTX

    dtx_values = filtered_df["dtx"].unique()

    for idx, out_fmt in enumerate(out_fmts):

        plot_loudness_by_bandwidth(df, in_fmt, out_fmt, out_dir)

        progressbar_update(idx + 1, len(out_fmts), width=50)

    print()