src/third_party/MusicSourceSeparationTraining/scripts/moises_to_musdb.py

import argparse
import os
import shutil
import time
from multiprocessing import Pool
from typing import Dict, List, Optional, Tuple, Union

import numpy as np
import soundfile as sf
from tqdm import tqdm


def combine_audio_files(files: List[str]) -> Tuple[np.ndarray, int]:
    """
    Combines multiple audio files into one by overlaying them.
    Parameters:
    - files (List[str]): List of file paths to be combined.
    Returns:
    - Tuple[np.ndarray, int]: A Tuple containing the combined audio data array and sample rate.
    """
    combined_data, sample_rate = sf.read(files[0])

    for file in files[1:]:
        data, sr = sf.read(file)
        if len(data) > len(combined_data):
            combined_data = np.pad(
                combined_data, ((0, len(data) - len(combined_data)), (0, 0)), "constant"
            )
        elif len(combined_data) > len(data):
            data = np.pad(
                data, ((0, len(combined_data) - len(data)), (0, 0)), "constant"
            )
        combined_data += data

    return combined_data, sample_rate


# Finds all .wav files located in folders that do not contain specified categories in their folder names
def files_to_categories(src_folder: str, categories: List[str]) -> Dict[str, List[str]]:
    """
    Finds all .wav files located in folders that do not contain specified categories
    in their folder names, within the given src_folder directory.
    Parameters:
    - src_folder (str): Path to the main directory containing subdirectories with files.
    - categories (List[str]): Keywords that should not be part of the folder's name.
    Returns:
    - Dict[str, List[str]]: A Dict with keys as categories, values as lists of paths to .wav files found.
    """
    files = {category: [] for category in categories + ["other"]}

    for folder in os.listdir(src_folder):
        folder_path = os.path.join(src_folder, folder)
        if os.path.isdir(folder_path):
            if folder.lower() in categories:
                stem = folder.lower()
            else:
                stem = "other"
            for f in os.listdir(folder_path):
                if f.endswith(".wav"):
                    files[stem].append(os.path.join(folder_path, f))
    return files


# Processes a folder containing audio tracks, copying and combining necessary files into the target structure
def process_folder(
    src_folder: str, dest_folder: str, stems: List[str], trim: bool = False
) -> None:
    """
    Processes a folder containing audio tracks, copying and combining necessary files into the target structure.
    Parameters:
    - src_folder (str): Path to the source folder of MoisesDB.
    - dest_folder (str): Path to the target folder for MUSDB18.
    - stems (List[str]): List of stem categories to process.
    - trim (bool): If True, trim all stems to the length of the shortest one.
    """

    if not os.path.exists(dest_folder):
        os.makedirs(dest_folder)

    categories = stems

    if trim:
        # First pass: load all stems and find the minimum length
        stem_data = {}
        sample_rate = None
        min_length = float("inf")
        all_files = files_to_categories(src_folder, categories)

        # Using tqdm to display progress for categories
        for category in tqdm(
            categories, desc=f"Processing categories in {os.path.basename(src_folder)}"
        ):
            files = all_files[category]
            if files:
                if len(files) > 1:
                    combined_data, sr = combine_audio_files(files)
                else:
                    combined_data, sr = sf.read(files[0])

                stem_data[category] = combined_data
                sample_rate = sr
                min_length = min(min_length, len(combined_data))

        # Process 'other' files
        other_files = all_files["other"]
        if other_files:
            other_combined_data, sr = combine_audio_files(other_files)
            stem_data["other"] = other_combined_data
            sample_rate = sr
            min_length = min(min_length, len(other_combined_data))

        # If no stems were found, set a default sample rate
        if sample_rate is None:
            sample_rate = 44100
            min_length = 0

        # Second pass: trim all stems to min_length and save
        for category in categories:
            if category in stem_data:
                trimmed_data = stem_data[category][:min_length]
                sf.write(
                    os.path.join(dest_folder, f"{category}.wav"),
                    trimmed_data,
                    sample_rate,
                )
            else:
                # Create silence with min_length
                silence = np.zeros((min_length, 2), dtype=np.float32)
                sf.write(
                    os.path.join(dest_folder, f"{category}.wav"), silence, sample_rate
                )

        # Save 'other' stem
        if "other" in stem_data:
            trimmed_other = stem_data["other"][:min_length]
            sf.write(os.path.join(dest_folder, "other.wav"), trimmed_other, sample_rate)
        else:
            silence = np.zeros((min_length, 2), dtype=np.float32)
            sf.write(os.path.join(dest_folder, "other.wav"), silence, sample_rate)

        # Create mixture.wav from all files, then trim to min_length
        all_files_list = [file for sublist in all_files.values() for file in sublist]
        if all_files_list:
            mixture_data, sample_rate = combine_audio_files(all_files_list)
            mixture_data = mixture_data[:min_length]
            sf.write(
                os.path.join(dest_folder, "mixture.wav"), mixture_data, sample_rate
            )
        else:
            # If no files at all, create silent mixture
            silence = np.zeros((min_length, 2), dtype=np.float32)
            sf.write(os.path.join(dest_folder, "mixture.wav"), silence, sample_rate)
    else:
        # Original behavior: If the required stem does not exist in the source folder (src_folder),
        # we add silence instead of the file with the same duration as the standard file.
        problem_categories = []
        duration = 0
        all_files = files_to_categories(src_folder, categories)

        # Using tqdm to display progress for categories
        for category in tqdm(
            categories, desc=f"Processing categories in {os.path.basename(src_folder)}"
        ):
            files = all_files[category]
            if files:
                if len(files) > 1:
                    combined_data, sample_rate = combine_audio_files(files)
                else:
                    combined_data, sample_rate = sf.read(files[0])

                sf.write(
                    os.path.join(dest_folder, f"{category}.wav"),
                    combined_data,
                    sample_rate,
                )
                duration = max(duration, len(combined_data) / sample_rate)
            else:
                problem_categories.append(category)

        other_files = all_files["other"]
        if other_files:
            other_combined_data, sample_rate = combine_audio_files(other_files)
            sf.write(
                os.path.join(dest_folder, "other.wav"), other_combined_data, sample_rate
            )
        else:
            problem_categories.append("other")

        for category in problem_categories:
            silence = np.zeros((int(duration * sample_rate), 2), dtype=np.float32)
            sf.write(os.path.join(dest_folder, f"{category}.wav"), silence, sample_rate)
        # mixture.wav
        all_files_list = [file for sublist in all_files.values() for file in sublist]
        mixture_data, sample_rate = combine_audio_files(all_files_list)
        sf.write(os.path.join(dest_folder, "mixture.wav"), mixture_data, sample_rate)


# Wrapper function for 'process_folder' that unpacks the arguments
def process_folder_wrapper(args: Tuple[str, str, List[str], bool]) -> None:
    """
    A wrapper function for 'process_folder' that unpacks the arguments.
    Parameters:
    - args (Tuple[str, str, List[str], bool]): A Tuple containing the source folder, destination folder paths, stems, and trim flag.
    """
    src_folder, dest_folder, stems, trim = args
    return process_folder(src_folder, dest_folder, stems, trim)


# Converts MoisesDB dataset to MUSDB18 format for a specified number of folders
def convert_dataset(
    src_root: str,
    dest_root: str,
    stems: List[str],
    max_folders: int = 240,
    num_workers: int = 4,
    trim: bool = False,
) -> None:
    """
    Converts MoisesDB dataset to MUSDB18 format for a specified number of folders.
    Parameters:
    - src_root (str): Root directory of the MoisesDB dataset.
    - dest_root (str): Root directory where the new dataset will be saved.
    - max_folders (int): Maximum number of folders to process.
    - num_workers (int): Number of parallel workers for processing.
    - trim (bool): If True, trim all stems to the length of the shortest one.
    """
    folders_to_process = []
    for folder in os.listdir(src_root):
        if len(folders_to_process) >= max_folders:
            break

        src_folder = os.path.join(src_root, folder)
        dest_folder = os.path.join(dest_root, folder)

        if os.path.isdir(src_folder):
            folders_to_process.append((src_folder, dest_folder, stems, trim))
        else:
            print(f"Skip {src_folder} — not dir")

    with Pool(num_workers) as pool:
        pool.map(process_folder_wrapper, folders_to_process)


# Count number of subfolders in a folder
def count_folders_in_folder(args_to_func) -> Dict[str, int]:
    """
    Counts the number of subfolders inside a given folder.

    Parameters:
    - folder_path (str): Path to the folder where the count is needed.

    Returns:
    - Dict[str, int]: A dictionary with folder paths as keys and subfolder counts as values.
    """
    folder_count = 0
    folder_path, stems = args_to_func
    if os.path.isdir(folder_path):
        # Count subfolders in stems
        folder_count = len(
            [
                f
                for f in os.listdir(folder_path)
                if os.path.isdir(os.path.join(folder_path, f)) and f in stems
            ]
        )
        # For other.wav
        if any(
            os.path.isdir(os.path.join(folder_path, f)) and f not in stems
            for f in os.listdir(folder_path)
        ):
            folder_count += 1

    return {folder_path: folder_count}


# Parallel count of subfolders in each folder inside src_folder
def count_folders_parallel(
    src_folder: str, stems, num_workers: int = 4
) -> Dict[str, int]:
    """
    Parallelly counts the number of subfolders in each folder inside src_folder.

    Parameters:
    - src_folder (str): Root folder containing subfolders to count.

    Returns:
    - Dict[str, int]: Dictionary with folder paths as keys and subfolder counts as values.
    """
    # Get list of all folders inside src_folder
    folders_to_process = [
        os.path.join(src_folder, folder)
        for folder in os.listdir(src_folder)
        if os.path.isdir(os.path.join(src_folder, folder))
    ]

    args_to_func = [(folder, stems) for folder in folders_to_process]

    # Parallelly process each folder using pool.map
    with Pool(num_workers) as pool:
        results = pool.map(count_folders_in_folder, args_to_func)

    # Merge results from different processes
    merged_counts = {}
    for result in results:
        merged_counts.update(result)

    return merged_counts


def parse_args(dict_args: Union[Dict, None]) -> argparse.Namespace:
    """
    Parse command-line arguments for configuring the model, dataset, and training parameters.

    Args:
        dict_args: Dict of command-line arguments. If None, arguments will be parsed from sys.argv.

    Returns:
        Namespace object containing parsed arguments and their values.
    """
    parser = argparse.ArgumentParser(
        description="Copy mixture files from VALID_DIR to INFERENCE_DIR"
    )
    parser.add_argument(
        "--src_dir",
        type=str,
        required=True,
        help="Source directory with MoisesDB tracks",
    )
    parser.add_argument(
        "--dest_dir",
        type=str,
        required=True,
        help="Directory to save tracks in MUSDB18",
    )
    parser.add_argument(
        "--num_workers", type=int, default=os.cpu_count(), help="Num of processors"
    )
    parser.add_argument(
        "--max_folders", type=int, default=240, help="Num of folders to use"
    )
    parser.add_argument(
        "--create_valid", action="store_true", help="Create valid folders or not"
    )
    parser.add_argument(
        "--valid_dir", type=str, default=r"\valid", help="Directory for valid"
    )
    parser.add_argument(
        "--valid_size", type=int, default=10, help="Num of folders to use in valitd"
    )
    parser.add_argument(
        "--stems",
        nargs="+",
        type=str,
        default=["bass", "drums", "vocals"],
        choices=[
            "drums",
            "guitar",
            "vocals",
            "bass",
            "other_keys",
            "piano",
            "percussion",
            "bowed_strings",
            "wind",
            "other_plucked",
        ],
        help="List of stems to use.",
    )
    parser.add_argument(
        "--mixture_name", type=str, default="mixture.wav", help="Name of mixture tracks"
    )
    parser.add_argument(
        "--trim",
        action="store_true",
        help="Trim all stems to the length of the shortest one",
    )

    if dict_args is not None:
        args = parser.parse_args([])
        args_dict = vars(args)
        args_dict.update(dict_args)
        args = argparse.Namespace(**args_dict)
    else:
        args = parser.parse_args()

    return args


def main(args: Optional[argparse.Namespace] = None) -> None:
    start = time.time()

    args = parse_args(args)

    source_directory = args.src_dir
    destination_directory = args.dest_dir
    num_workers = args.num_workers
    stems = args.stems
    max_folders = args.max_folders
    trim = args.trim

    print(f"num_workers: {num_workers}, categories: {stems + ['other']}, trim: {trim}")

    convert_dataset(
        source_directory,
        destination_directory,
        stems,
        max_folders=max_folders,
        num_workers=num_workers,
        trim=trim,
    )

    print(
        f"All {max_folders} files have been processed, time: {time.time() - start:.2f} sec"
    )

    if args.create_valid:
        # Count folders in the MoisesDB dataset
        result = count_folders_parallel(source_directory, stems)
        result = dict(sorted(result.items(), key=lambda item: item[1], reverse=True))
        # valid_size = min(args.valid_size, max_folders)
        valid_size = 10
        list_folders = list(result.keys())
        print(f"Top {valid_size} folders:")
        for track in list(result.items())[:valid_size]:
            print(track)

        valid_folder = args.valid_dir
        train_tracks_folder = destination_directory

        # Create valid folder if not exists
        if not os.path.exists(valid_folder):
            os.makedirs(valid_folder)

        num_val = 0

        # Copy folders from train_tracks to valid folder
        for folder in list_folders:
            folder_name = os.path.basename(folder)  # Get folder name

            # Form the path to the folder in train_tracks
            source_folder = os.path.join(train_tracks_folder, folder_name)

            # If the folder exists in train_tracks, copy it to valid
            if os.path.exists(source_folder):
                destination = os.path.join(valid_folder, folder_name)
                shutil.copytree(source_folder, destination)
                shutil.rmtree(source_folder)
                num_val += 1
                print(f"Folder: {folder}, num_stems: {result[folder]}")
                if num_val >= valid_size:
                    break
            else:
                print(f"Folder {source_folder} not found.")

    print("The end!")


if __name__ == "__main__":
    main(None)