ttsds 0.0.1__py3-none-any.whl

ttsds/__about__.py

@@ -0,0 +1,4 @@
+# SPDX-FileCopyrightText: 2024-present Christoph Minixhofer <christoph.minixhofer@gmail.com>
+#
+# SPDX-License-Identifier: MIT
+__version__ = "0.0.1"

ttsds/__init__.py

@@ -0,0 +1,345 @@
+from typing import List, Optional
+import importlib.resources
+from time import time
+from pathlib import Path
+import pickle
+import gzip
+import requests
+
+import pandas as pd
+from transformers import logging
+import numpy as np
+from sklearn.decomposition import PCA
+
+from ttsds.benchmarks.environment.voicefixer import VoiceFixerBenchmark
+from ttsds.benchmarks.environment.wada_snr import WadaSNRBenchmark
+from ttsds.benchmarks.general.hubert import HubertBenchmark
+from ttsds.benchmarks.general.wav2vec2 import Wav2Vec2Benchmark
+from ttsds.benchmarks.general.wavlm import WavLMBenchmark
+from ttsds.benchmarks.intelligibility.w2v2_wer import Wav2Vec2WERBenchmark
+from ttsds.benchmarks.intelligibility.whisper_wer import WhisperWERBenchmark
+from ttsds.benchmarks.prosody.mpm import MPMBenchmark
+from ttsds.benchmarks.prosody.pitch import PitchBenchmark
+from ttsds.benchmarks.prosody.hubert_token import HubertTokenBenchmark
+from ttsds.benchmarks.speaker.wespeaker import WeSpeakerBenchmark
+from ttsds.benchmarks.speaker.dvector import DVectorBenchmark
+from ttsds.benchmarks.benchmark import BenchmarkCategory, BenchmarkDimension
+from ttsds.util.dataset import Dataset, TarDataset, DataDistribution, DEFAULT_BENCHMARKS
+
+# we do this to avoid "some weights of the model checkpoint at ... were not used when initializing" warnings
+logging.set_verbosity_error()
+
+
+benchmark_dict = {
+    "hubert": HubertBenchmark,
+    "wav2vec2": Wav2Vec2Benchmark,
+    "wavlm": WavLMBenchmark,
+    "wav2vec2_wer": Wav2Vec2WERBenchmark,
+    "whisper_wer": WhisperWERBenchmark,
+    "mpm": MPMBenchmark,
+    "pitch": PitchBenchmark,
+    "wespeaker": WeSpeakerBenchmark,
+    "dvector": DVectorBenchmark,
+    "hubert_token": HubertTokenBenchmark,
+    "voicefixer": VoiceFixerBenchmark,
+    "wada_snr": WadaSNRBenchmark,
+}
+
+with importlib.resources.path("ttsds", "data") as data_path:
+    # if they don't exist, download from github
+    for noise_name in [
+        "esc50",
+        "all_ones",
+        "all_zeros",
+        "normal_distribution",
+        "uniform_distribution",
+    ]:
+        if not Path(f"{data_path}/noise_{noise_name}.pkl.gz").exists():
+            print(f"Downloading noise_{noise_name}.pkl.gz")
+            url = f"https://github.com/ttsds/ttsds/raw/main/src/ttsds/data/noise_{noise_name}.pkl.gz"
+            r = requests.get(url)
+            with open(f"{data_path}/noise_{noise_name}.pkl.gz", "wb") as f:
+                f.write(r.content)
+
+    for speech_name in [
+        "blizzard2008",
+        "blizzard2013",
+        "common_voice",
+        "libritts_test",
+        "libritts_r_test",
+        "lj_speech",
+        "vctk",
+    ]:
+        if not Path(f"{data_path}/reference_speech_{speech_name}.pkl.gz").exists():
+            print(f"Downloading reference_speech_{speech_name}.pkl.gz")
+            url = f"https://github.com/ttsds/ttsds/raw/main/src/ttsds/data/reference_speech_{speech_name}.pkl.gz"
+            r = requests.get(url)
+            with open(f"{data_path}/reference_speech_{speech_name}.pkl.gz", "wb") as f:
+                f.write(r.content)
+
+
+    # check if the reference and noise distributions are already saved
+    if not Path(f"{data_path}/reference_speech_blizzard2008.pkl.gz").exists():
+        print("Creating reference distributions")
+        ref_benchmark_dict = {
+            k: v() for k, v in benchmark_dict.items()
+        }
+        REFERENCE_DISTS = [
+            DataDistribution(
+                TarDataset(data_path / "original" / f"speech_{name}.tar.gz"),
+                ref_benchmark_dict,
+                benchmarks=DEFAULT_BENCHMARKS,
+                name=f"speech_{name}",
+            )
+            for name in [
+                "blizzard2008",
+                "blizzard2013",
+                "common_voice",
+                "libritts_test",
+                "libritts_r_test",
+                "lj_speech",
+                "vctk",
+            ]
+        ]
+        # save the reference distributions
+        for dist in REFERENCE_DISTS:
+            dist.to_pickle(f"{data_path}/reference_{dist.name}.pkl.gz")
+
+    REFERENCE_DISTS = [
+        DataDistribution.from_pickle(f"{data_path}/reference_{name}.pkl.gz")
+        for name in [
+            "speech_blizzard2008",
+            "speech_blizzard2013",
+            "speech_common_voice",
+            "speech_libritts_test",
+            "speech_libritts_r_test",
+            "speech_lj_speech",
+            "speech_vctk",
+        ]
+    ]
+
+    if not Path(f"{data_path}/noise_esc50.pkl.gz").exists():
+        print("Creating noise distributions")
+        ref_benchmark_dict = {
+            k: v() for k, v in benchmark_dict.items()
+        }
+        NOISE_DISTS = [
+            DataDistribution(
+                TarDataset(data_path / "original" / f"noise_{name}.tar.gz"),
+                ref_benchmark_dict,
+                benchmarks=DEFAULT_BENCHMARKS,
+                name=name,
+            )
+            for name in [
+                "esc50",
+                "all_ones",
+                "all_zeros",
+                "normal_distribution",
+                "uniform_distribution",
+            ]
+        ]
+        # save the noise distributions
+        for dist in NOISE_DISTS:
+            dist.to_pickle(f"{data_path}/noise_{dist.name}.pkl.gz")
+
+    NOISE_DISTS = [
+        DataDistribution.from_pickle(f"{data_path}/noise_{name}.pkl.gz")
+        for name in [
+            "esc50",
+            "all_ones",
+            "all_zeros",
+            "normal_distribution",
+            "uniform_distribution",
+        ]
+    ]
+
+
+class BenchmarkSuite:
+    def __init__(
+        self,
+        datasets: List[Dataset],
+        benchmarks: List[str] = DEFAULT_BENCHMARKS,
+        print_results: bool = True,
+        skip_errors: bool = False,
+        noise_distributions: List[DataDistribution] = NOISE_DISTS,
+        reference_distributions: List[DataDistribution] = REFERENCE_DISTS,
+        write_to_file: str = None,
+    ):
+        self.benchmarks = benchmarks
+        self.benchmark_objects = [benchmark_dict[benchmark]() for benchmark in benchmarks]
+        # sort by category and then by name
+        self.benchmark_objects = sorted(
+            self.benchmark_objects, key=lambda x: (x.category.value, x.name)
+        )
+        self.datasets = datasets
+        self.datasets = sorted(self.datasets, key=lambda x: x.name)
+        self.database = pd.DataFrame(
+            columns=[
+                "benchmark_name",
+                "benchmark_category",
+                "dataset",
+                "score",
+                "ci",
+                "time_taken",
+                "noise_dataset",
+                "reference_dataset",
+            ]
+        )
+        self.print_results = print_results
+        self.skip_errors = skip_errors
+        self.noise_distributions = noise_distributions
+        self.reference_distributions = reference_distributions
+        self.write_to_file = write_to_file
+        if Path(write_to_file).exists():
+            self.database = pd.read_csv(write_to_file, index_col=0)
+            self.database = self.database.reset_index()
+
+    def run(self) -> pd.DataFrame:
+        for benchmark in self.benchmark_objects:
+            for dataset in self.datasets:
+                # empty lines for better readability
+                print("\n")
+                print(f"{'='*80}")
+                print(f"Benchmark Category: {benchmark.category.name}")
+                print(f"Running {benchmark.name} on {dataset.root_dir}")
+                try:
+                    # check if it's in the database
+                    if (
+                        (self.database["benchmark_name"] == benchmark.name)
+                        & (self.database["dataset"] == dataset.name)
+                    ).any():
+                        print(
+                            f"Skipping {benchmark.name} on {dataset.name} as it's already in the database"
+                        )
+                        continue
+                    start = time()
+                    if "WER".lower() in benchmark.name.lower():
+                        print([
+                            x.get_distribution(benchmark.key) for x in self.reference_distributions
+                        ])
+                        print([
+                            x.get_distribution(benchmark.key) for x in self.noise_distributions
+                        ])
+                        print(benchmark.get_distribution(dataset))
+                    score = benchmark.compute_score(
+                        dataset, self.reference_distributions, self.noise_distributions
+                    )
+                    time_taken = time() - start
+                except Exception as e:
+                    if self.skip_errors:
+                        print(f"Error: {e}")
+                        score = (np.nan, np.nan)
+                        time_taken = np.nan
+                    else:
+                        raise e
+                result = {
+                    "benchmark_name": [benchmark.name],
+                    "benchmark_category": [benchmark.category.value],
+                    "dataset": [dataset.name],
+                    "score": [score[0]],
+                    "ci": [score[1]],
+                    "time_taken": [time_taken],
+                    "noise_dataset": [score[2][0]],
+                    "reference_dataset": [score[2][1]],
+                }
+                if self.print_results:
+                    print(result)
+                self.database = pd.concat(
+                    [
+                        self.database,
+                        pd.DataFrame(result),
+                    ],
+                    ignore_index=True,
+                )
+                if self.write_to_file is not None:
+                    self.database["score"] = self.database["score"].astype(float)
+                    self.database = self.database.sort_values(
+                        ["benchmark_category", "benchmark_name", "score"],
+                        ascending=[True, True, False],
+                    )
+                    self.database.to_csv(self.write_to_file, index=False)
+        return self.database
+
+    @staticmethod
+    def aggregate_df(df: pd.DataFrame) -> pd.DataFrame:
+        def concat_text(x):
+            return ", ".join(x)
+
+        df["benchmark_category"] = df["benchmark_category"].apply(
+            lambda x: BenchmarkCategory(x).name
+        )
+        df = (
+            df.groupby(
+                [
+                    "benchmark_category",
+                    "dataset",
+                ]
+            )
+            .agg(
+                {
+                    "score": ["mean"],
+                    "ci": ["mean"],
+                    "time_taken": ["mean"],
+                    "noise_dataset": [concat_text],
+                    "reference_dataset": [concat_text],
+                    "benchmark_name": [concat_text],
+                }
+            )
+            .reset_index()
+        )
+        # remove multiindex
+        df.columns = [x[0] for x in df.columns.ravel()]
+        # drop the benchmark_name column
+        df = df.drop("benchmark_name", axis=1)
+        # replace benchmark_category number with string
+        return df
+
+    def get_aggregated_results(self) -> pd.DataFrame:
+        df = self.database.copy()
+        return BenchmarkSuite.aggregate_df(df)
+
+    def get_benchmark_distribution(
+        self,
+        benchmark_name: str,
+        dataset_name: str,
+        pca_components: Optional[int] = None,
+    ) -> dict:
+        benchmark = [x for x in self.benchmark_objects if x.name == benchmark_name][0]
+        dataset = [x for x in self.datasets if x.name == dataset_name][0]
+        closest_noise = self.database[
+            (self.database["benchmark_name"] == benchmark_name)
+            & (self.database["dataset"] == dataset_name)
+        ]["noise_dataset"].values[0]
+        closest_noise = [
+            x for x in self.noise_distributions if x.name == closest_noise
+        ][0]
+        other_noise = [
+            x for x in self.noise_distributions if x.name != closest_noise.name
+        ][0]
+        closest_reference = self.database[
+            (self.database["benchmark_name"] == benchmark_name)
+            & (self.database["dataset"] == dataset_name)
+        ]["reference_dataset"].values[0]
+        closest_reference = [
+            x for x in self.reference_distributions if x.name == closest_reference
+        ][0]
+        other_reference = [
+            x for x in self.reference_distributions if x.name != closest_reference.name
+        ][0]
+        result = {
+            "benchmark_distribution": benchmark.get_distribution(dataset),
+            "noise_distribution": benchmark.get_distribution(closest_noise),
+            "reference_distribution": benchmark.get_distribution(closest_reference),
+            "other_noise_distribution": benchmark.get_distribution(other_noise),
+            "other_reference_distribution": benchmark.get_distribution(other_reference),
+        }
+        if pca_components is not None:
+            pca = PCA(n_components=pca_components)
+            # fit on all except the benchmark distribution
+            pca.fit(
+                np.vstack(
+                    [v for k, v in result.items() if k != "benchmark_distribution"]
+                )
+            )
+            result = {k: pca.transform(v) for k, v in result.items()}
+        return result

ttsds/benchmarks/benchmark.py

@@ -0,0 +1,182 @@
+"""
+This file contains the Benchmark abstract class.
+"""
+
+from abc import ABC, abstractmethod
+from enum import Enum
+import hashlib
+import importlib.resources
+import json
+from typing import List, Union
+from functools import lru_cache
+from pathlib import Path
+
+import numpy as np
+
+from ttsds.util.dataset import Dataset, DataDistribution
+from ttsds.util.cache import cache, load_cache, check_cache, hash_md5
+from ttsds.util.distances import wasserstein_distance, frechet_distance
+
+
+class BenchmarkCategory(Enum):
+    """
+    Enum class for the different categories of benchmarks.
+    """
+
+    OVERALL = 1
+    PROSODY = 2
+    ENVIRONMENT = 3
+    SPEAKER = 4
+    PHONETICS = 5
+    INTELLIGIBILITY = 6
+    TRAINABILITY = 7
+    EXTERNAL = 8
+
+
+class BenchmarkDimension(Enum):
+    """
+    Enum class for the different dimensions of benchmarks.
+    """
+
+    ONE_DIMENSIONAL = 1
+    N_DIMENSIONAL = 2
+
+
+class Benchmark(ABC):
+    """
+    Abstract class for a benchmark.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        category: BenchmarkCategory,
+        dimension: BenchmarkDimension,
+        description: str,
+        **kwargs,
+    ):
+        self.name = name
+        self.key = name.lower().replace(" ", "_")
+        self.category = category
+        self.dimension = dimension
+        self.description = description
+        self.kwargs = kwargs
+
+    def get_distribution(self, dataset: Union[Dataset, DataDistribution]) -> np.ndarray:
+        """
+        Abstract method to get the distribution of the benchmark.
+        If the benchmark is one-dimensional, the method should return a
+        numpy array with the values of the benchmark for each sample in the dataset.
+        If the benchmark is n-dimensional, the method should return a numpy array
+        with the values of the benchmark for each sample in the dataset, where each
+        row corresponds to a sample and each column corresponds to a dimension of the benchmark.
+        """
+        ds_hash = hash_md5(dataset)
+        benchmark_hash = hash_md5(self)
+        cache_name = f"benchmarks/{self.name}/{ds_hash}_{benchmark_hash}"
+        if check_cache(cache_name):
+            return load_cache(cache_name)
+        if check_cache(cache_name + "_mu") and check_cache(cache_name + "_sig"):
+            mu = load_cache(cache_name + "_mu")
+            sig = load_cache(cache_name + "_sig")
+            return (mu, sig)
+        if isinstance(dataset, DataDistribution) and self.dimension == BenchmarkDimension.N_DIMENSIONAL:
+            mu, sig = dataset.get_distribution(self.key)
+            cache(mu, cache_name + "_mu")
+            cache(sig, cache_name + "_sig")
+            return (mu, sig)
+        elif isinstance(dataset, DataDistribution) and self.dimension == BenchmarkDimension.ONE_DIMENSIONAL:
+            distribution = dataset.get_distribution(self.key)
+            cache(distribution, cache_name)
+            return distribution
+        distribution = self._get_distribution(dataset)
+        cache(distribution, cache_name)
+        return distribution
+
+    @abstractmethod
+    def _get_distribution(self, dataset: Dataset) -> np.ndarray:
+        """
+        Abstract method to get the distribution of the benchmark.
+        """
+        raise NotImplementedError
+
+    def __str__(self) -> str:
+        return f"{self.category.name}/{self.name}"
+
+    def __repr__(self):
+        return f"{self.category.name}/{self.name}"
+
+    def __hash__(self) -> int:
+        h = hashlib.md5()
+        h.update(self.name.encode())
+        h.update(self.category.name.encode())
+        h.update(self.dimension.name.encode())
+        h.update(self.description.encode())
+        # convert the kwargs to strings
+        kwargs_str = {
+            k: str(v) if not isinstance(v, dict) else json.dumps(v, sort_keys=True)
+            for k, v in self.kwargs.items()
+        }
+        h.update(json.dumps(kwargs_str, sort_keys=True).encode())
+        return int(h.hexdigest(), 16)
+
+    @lru_cache(maxsize=None)
+    def compute_distance(
+        self,
+        one_dataset: Union[Dataset, DataDistribution],
+        other_dataset: Union[Dataset, DataDistribution],
+    ) -> float:
+        """
+        Compute the distance between the distributions of the benchmark in two datasets.
+        """
+        one_distribution = self.get_distribution(one_dataset)
+        other_distribution = self.get_distribution(other_dataset)
+        if self.dimension == BenchmarkDimension.ONE_DIMENSIONAL:
+            return wasserstein_distance(one_distribution, other_distribution)
+        elif self.dimension == BenchmarkDimension.N_DIMENSIONAL:
+            return frechet_distance(one_distribution, other_distribution)
+        else:
+            raise ValueError("Invalid benchmark dimension")
+
+    def compute_score(
+        self,
+        dataset: Dataset,
+        reference_datasets: List[Dataset],
+        noise_datasets: List[Dataset],
+    ) -> float:
+        """
+        Compute the score of the benchmark on a dataset.
+        """
+        noise_scores = []
+        for noise_ds in noise_datasets:
+            score = self.compute_distance(noise_ds, dataset)
+            noise_scores.append(score)
+        noise_scores = np.array(noise_scores)
+
+        dataset_scores = []
+        for ref_ds in reference_datasets:
+            score = self.compute_distance(ref_ds, dataset)
+            dataset_scores.append(score)
+        dataset_scores = np.array(dataset_scores)
+
+        closest_noise_idx = np.argmin(noise_scores)
+        closest_dataset_idx = np.argmin(dataset_scores)
+
+        print(f"Closest noise dataset: {noise_datasets[closest_noise_idx].name}")
+        print(
+            f"Closest reference dataset: {reference_datasets[closest_dataset_idx].name}"
+        )
+
+        noise_score = np.min(noise_scores)
+        dataset_score = np.min(dataset_scores)
+        combined_score = dataset_score + noise_score
+        score = (noise_score / combined_score) * 100
+        # TODO: compute confidence interval
+        return (
+            score,
+            1.0,
+            (
+                noise_datasets[closest_noise_idx].name,
+                reference_datasets[closest_dataset_idx].name,
+            ),
+        )

ttsds/benchmarks/environment/voicefixer.py

@@ -0,0 +1,80 @@
+import tempfile
+
+from pesq import pesq
+from voicefixer import VoiceFixer
+from simple_hifigan import Synthesiser
+import numpy as np
+import soundfile as sf
+import librosa
+from tqdm import tqdm
+
+from ttsds.benchmarks.benchmark import Benchmark, BenchmarkCategory, BenchmarkDimension
+from ttsds.util.dataset import Dataset
+
+
+class VoiceFixerBenchmark(Benchmark):
+    """
+    Benchmark class for the VoiceFixer benchmark.
+    """
+
+    def __init__(
+        self,
+    ):
+        super().__init__(
+            name="VoiceFixer",
+            category=BenchmarkCategory.ENVIRONMENT,
+            dimension=BenchmarkDimension.ONE_DIMENSIONAL,
+            description="The phone counts of VoiceFixer.",
+        )
+        self.model = VoiceFixer()
+        self.synthesiser = Synthesiser()
+
+    def _get_distribution(self, dataset: Dataset) -> np.ndarray:
+        """
+        Compute the Word Error Rate (WER) distribution of the VoiceFixer model.
+
+        Args:
+            dataset (Dataset): The dataset to compute the WER on.
+
+        Returns:
+            float: The Word Error Rate (WER) distribution of the VoiceFixer model.
+        """
+        mel_diffs = []
+        for wav, _ in tqdm(dataset, desc=f"computing noise for {self.name}"):
+            if dataset.sample_rate != 16000:
+                wav = librosa.resample(
+                    wav, orig_sr=dataset.sample_rate, target_sr=16000
+                )
+            with tempfile.NamedTemporaryFile(suffix=".wav") as f:
+                # take random 2 seconds
+                if len(wav) > 32000:
+                    start = np.random.randint(0, len(wav) - 32000)
+                    wav = wav[start : start + 32000]
+                sf.write(f.name, wav, 16000)
+                with tempfile.NamedTemporaryFile(suffix=".wav") as f_out:
+                    self.model.restore(f.name, f_out.name)
+                    wav_out, _ = librosa.load(f_out.name, sr=16000)
+            wav = wav / (np.max(np.abs(wav)) + 1e-5)
+            wav_out = wav_out / np.max(np.abs(wav_out))
+            mel = self.synthesiser.wav_to_mel(wav, 16000)[0].T
+            mel_out = self.synthesiser.wav_to_mel(wav_out, 16000)[0].T
+            if mel_out.shape[0] > mel.shape[0]:
+                mel_out = mel_out[: mel.shape[0]]
+            elif mel_out.shape[0] < mel.shape[0]:
+                mel = mel[: mel_out.shape[0]]
+            mel_diff = mel_out
+            # check if there is any nan
+            if np.isnan(mel_diff).any():
+                print("nan found, skip")
+                continue
+            # convert back to wav
+            mel_diff = self.synthesiser(mel_diff.T)[0]
+            mel_diff = mel_diff / np.max(np.abs(mel_diff) + 1e-5)
+            mel_diff = librosa.resample(mel_diff, orig_sr=22050, target_sr=16000)
+            # calculate the difference
+            try:
+                mel_diff = pesq(16000, wav, mel_diff, "wb")
+                mel_diffs.append(mel_diff)
+            except:
+                mel_diffs.append(0)
+        return np.array(mel_diffs)