Upload 5 files

metrics/UTMOS.py

@@ -0,0 +1,223 @@
+import os
+
+import fairseq
+import pytorch_lightning as pl
+import requests
+import torch
+import torch.nn as nn
+from tqdm import tqdm
+
+UTMOS_CKPT_URL = "https://huggingface.co/spaces/sarulab-speech/UTMOS-demo/resolve/main/epoch%3D3-step%3D7459.ckpt"
+WAV2VEC_URL = "https://huggingface.co/spaces/sarulab-speech/UTMOS-demo/resolve/main/wav2vec_small.pt"
+
+"""
+UTMOS score, automatic Mean Opinion Score (MOS) prediction system, 
+adapted from https://huggingface.co/spaces/sarulab-speech/UTMOS-demo
+"""
+
+
+class UTMOSScore:
+    """Predicting score for each audio clip."""
+
+    def __init__(self, device, ckpt_path="epoch=3-step=7459.ckpt"):
+        self.device = device
+        filepath = os.path.join(os.path.dirname(__file__), ckpt_path)
+        if not os.path.exists(filepath):
+            download_file(UTMOS_CKPT_URL, filepath)
+        self.model = BaselineLightningModule.load_from_checkpoint(filepath).eval().to(device)
+
+    def score(self, wavs: torch.tensor) -> torch.tensor:
+        """
+        Args:
+            wavs: audio waveform to be evaluated. When len(wavs) == 1 or 2,
+                the model processes the input as a single audio clip. The model
+                performs batch processing when len(wavs) == 3.
+        """
+        if len(wavs.shape) == 1:
+            out_wavs = wavs.unsqueeze(0).unsqueeze(0)
+        elif len(wavs.shape) == 2:
+            out_wavs = wavs.unsqueeze(0)
+        elif len(wavs.shape) == 3:
+            out_wavs = wavs
+        else:
+            raise ValueError("Dimension of input tensor needs to be <= 3.")
+        bs = out_wavs.shape[0]
+        batch = {
+            "wav": out_wavs,
+            "domains": torch.zeros(bs, dtype=torch.int).to(self.device),
+            "judge_id": torch.ones(bs, dtype=torch.int).to(self.device) * 288,
+        }
+        with torch.no_grad():
+            output = self.model(batch)
+
+        return output.mean(dim=1).squeeze(1).cpu().detach() * 2 + 3
+
+
+def download_file(url, filename):
+    """
+    Downloads a file from the given URL
+
+    Args:
+        url (str): The URL of the file to download.
+        filename (str): The name to save the file as.
+    """
+    print(f"Downloading file {filename}...")
+    response = requests.get(url, stream=True)
+    response.raise_for_status()
+
+    total_size_in_bytes = int(response.headers.get("content-length", 0))
+    progress_bar = tqdm(total=total_size_in_bytes, unit="iB", unit_scale=True)
+
+    with open(filename, "wb") as f:
+        for chunk in response.iter_content(chunk_size=8192):
+            progress_bar.update(len(chunk))
+            f.write(chunk)
+
+    progress_bar.close()
+
+
+def load_ssl_model(ckpt_path="wav2vec_small.pt"):
+    filepath = os.path.join(os.path.dirname(__file__), ckpt_path)
+    if not os.path.exists(filepath):
+        download_file(WAV2VEC_URL, filepath)
+    SSL_OUT_DIM = 768
+    model, cfg, task = fairseq.checkpoint_utils.load_model_ensemble_and_task([filepath])
+    ssl_model = model[0]
+    ssl_model.remove_pretraining_modules()
+    return SSL_model(ssl_model, SSL_OUT_DIM)
+
+
+class BaselineLightningModule(pl.LightningModule):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.construct_model()
+        self.save_hyperparameters()
+
+    def construct_model(self):
+        self.feature_extractors = nn.ModuleList(
+            [load_ssl_model(ckpt_path="wav2vec_small.pt"), DomainEmbedding(3, 128),]
+        )
+        output_dim = sum([feature_extractor.get_output_dim() for feature_extractor in self.feature_extractors])
+        output_layers = [LDConditioner(judge_dim=128, num_judges=3000, input_dim=output_dim)]
+        output_dim = output_layers[-1].get_output_dim()
+        output_layers.append(
+            Projection(hidden_dim=2048, activation=torch.nn.ReLU(), range_clipping=False, input_dim=output_dim)
+        )
+
+        self.output_layers = nn.ModuleList(output_layers)
+
+    def forward(self, inputs):
+        outputs = {}
+        for feature_extractor in self.feature_extractors:
+            outputs.update(feature_extractor(inputs))
+        x = outputs
+        for output_layer in self.output_layers:
+            x = output_layer(x, inputs)
+        return x
+
+
+class SSL_model(nn.Module):
+    def __init__(self, ssl_model, ssl_out_dim) -> None:
+        super(SSL_model, self).__init__()
+        self.ssl_model, self.ssl_out_dim = ssl_model, ssl_out_dim
+
+    def forward(self, batch):
+        wav = batch["wav"]
+        wav = wav.squeeze(1)  # [batches, audio_len]
+        res = self.ssl_model(wav, mask=False, features_only=True)
+        x = res["x"]
+        return {"ssl-feature": x}
+
+    def get_output_dim(self):
+        return self.ssl_out_dim
+
+
+class DomainEmbedding(nn.Module):
+    def __init__(self, n_domains, domain_dim) -> None:
+        super().__init__()
+        self.embedding = nn.Embedding(n_domains, domain_dim)
+        self.output_dim = domain_dim
+
+    def forward(self, batch):
+        return {"domain-feature": self.embedding(batch["domains"])}
+
+    def get_output_dim(self):
+        return self.output_dim
+
+
+class LDConditioner(nn.Module):
+    """
+    Conditions ssl output by listener embedding
+    """
+
+    def __init__(self, input_dim, judge_dim, num_judges=None):
+        super().__init__()
+        self.input_dim = input_dim
+        self.judge_dim = judge_dim
+        self.num_judges = num_judges
+        assert num_judges != None
+        self.judge_embedding = nn.Embedding(num_judges, self.judge_dim)
+        # concat [self.output_layer, phoneme features]
+
+        self.decoder_rnn = nn.LSTM(
+            input_size=self.input_dim + self.judge_dim,
+            hidden_size=512,
+            num_layers=1,
+            batch_first=True,
+            bidirectional=True,
+        )  # linear?
+        self.out_dim = self.decoder_rnn.hidden_size * 2
+
+    def get_output_dim(self):
+        return self.out_dim
+
+    def forward(self, x, batch):
+        judge_ids = batch["judge_id"]
+        if "phoneme-feature" in x.keys():
+            concatenated_feature = torch.cat(
+                (x["ssl-feature"], x["phoneme-feature"].unsqueeze(1).expand(-1, x["ssl-feature"].size(1), -1)), dim=2
+            )
+        else:
+            concatenated_feature = x["ssl-feature"]
+        if "domain-feature" in x.keys():
+            concatenated_feature = torch.cat(
+                (concatenated_feature, x["domain-feature"].unsqueeze(1).expand(-1, concatenated_feature.size(1), -1),),
+                dim=2,
+            )
+        if judge_ids != None:
+            concatenated_feature = torch.cat(
+                (
+                    concatenated_feature,
+                    self.judge_embedding(judge_ids).unsqueeze(1).expand(-1, concatenated_feature.size(1), -1),
+                ),
+                dim=2,
+            )
+            decoder_output, (h, c) = self.decoder_rnn(concatenated_feature)
+        return decoder_output
+
+
+class Projection(nn.Module):
+    def __init__(self, input_dim, hidden_dim, activation, range_clipping=False):
+        super(Projection, self).__init__()
+        self.range_clipping = range_clipping
+        output_dim = 1
+        if range_clipping:
+            self.proj = nn.Tanh()
+
+        self.net = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim), activation, nn.Dropout(0.3), nn.Linear(hidden_dim, output_dim),
+        )
+        self.output_dim = output_dim
+
+    def forward(self, x, batch):
+        output = self.net(x)
+
+        # range clipping
+        if self.range_clipping:
+            return self.proj(output) * 2.0 + 3
+        else:
+            return output
+
+    def get_output_dim(self):
+        return self.output_dim

metrics/infer.py

@@ -0,0 +1,116 @@
+# 测试各种指标
+import os
+import glob
+from UTMOS import UTMOSScore
+from periodicity import calculate_periodicity_metrics
+import torchaudio
+from pesq import pesq
+import numpy as np
+import torch
+import math
+from pystoi import stoi
+
+device=torch.device('cuda:0')
+
+# 如果是ljspeech，需要更换路径，更换数据读取逻辑，更换stoi的采样率
+
+def main():
+    prepath="./Result/Minicodec/infer/dac_nq4_all"
+    rawpath="./Data/libritts/test-clean"
+    # rawpath="./Data/LJSpeech-1.1/wavs"
+    preaudio = os.listdir(prepath)
+    rawaudio = []
+
+    UTMOS=UTMOSScore(device='cuda:0')
+    
+    # libritts
+    for i in range(len(preaudio)):
+        id1=preaudio[i].split('_')[0]
+        id2=preaudio[i].split('_')[1]
+        rawaudio.append(rawpath+"/"+id1+"/"+id2+"/"+preaudio[i])
+
+    # # ljspeech
+    # for i in range(len(preaudio)):
+    #     rawaudio.append(rawpath+"/"+preaudio[i])
+
+    utmos_sumgt=0
+    utmos_sumencodec=0
+    pesq_sumpre=0
+    f1score_sumpre=0
+    stoi_sumpre=[]
+    f1score_filt=0
+
+    for i in range(len(preaudio)):
+        print(i)
+        rawwav,rawwav_sr=torchaudio.load(rawaudio[i])
+        prewav,prewav_sr=torchaudio.load(prepath+"/"+preaudio[i])
+        # breakpoint()
+        rawwav=rawwav.to(device)
+        prewav=prewav.to(device)
+        # print(rawwav.size(),prewav.size())
+        # breakpoint()
+        rawwav_16k=torchaudio.functional.resample(rawwav, orig_freq=rawwav_sr, new_freq=16000)  #测试UTMOS的时候必须重采样
+        prewav_16k=torchaudio.functional.resample(prewav, orig_freq=prewav_sr, new_freq=16000)
+
+
+        # 1.UTMOS
+        print("****UTMOS_raw",i,UTMOS.score(rawwav_16k.unsqueeze(1))[0].item())
+        print("****UTMOS_encodec",i,UTMOS.score(prewav_16k.unsqueeze(1))[0].item())
+        utmos_sumgt+=UTMOS.score(rawwav_16k.unsqueeze(1))[0].item()
+        utmos_sumencodec+=UTMOS.score(prewav_16k.unsqueeze(1))[0].item()
+    
+
+        # breakpoint()
+
+        ## 2.PESQ  
+        min_len=min(rawwav_16k.size()[1],prewav_16k.size()[1])
+        rawwav_16k_pesq=rawwav_16k[:,:min_len].squeeze(0)
+        prewav_16k_pesq=prewav_16k[:,:min_len].squeeze(0)
+        pesq_score = pesq(16000, rawwav_16k_pesq.cpu().numpy(), prewav_16k_pesq.cpu().numpy(), "wb", on_error=1)
+        print("****PESQ",i,pesq_score)
+        pesq_sumpre+=pesq_score
+        # breakpoint()
+
+        ## 3.F1-score
+        min_len=min(rawwav_16k.size()[1],prewav_16k.size()[1])
+        rawwav_16k_f1score=rawwav_16k[:,:min_len]
+        prewav_16k_f1score=prewav_16k[:,:min_len]
+        periodicity_loss, pitch_loss, f1_score = calculate_periodicity_metrics(rawwav_16k_f1score,prewav_16k_f1score)
+        print("****f1",periodicity_loss, pitch_loss, f1_score,f1score_sumpre)
+        if(math.isnan(f1_score)):
+            f1score_filt+=1
+            print("*****",f1score_filt)
+        else:
+            f1score_sumpre+=f1_score
+        # breakpoint()
+
+
+        ## 4.STOI
+        # # 针对重采样的ljspeech
+        # rawwav_24k=torchaudio.functional.resample(rawwav, orig_freq=rawwav_sr, new_freq=24000)
+        # min_len=min(rawwav_24k.size()[1],prewav.size()[1])
+        # rawwav_stoi=rawwav_24k[:,:min_len].squeeze(0)
+        # prewav_stoi=prewav[:,:min_len].squeeze(0)
+        # tmp_stoi=stoi(rawwav_stoi.cpu(),prewav_stoi.cpu(),24000,extended=False)
+        # print("****stoi",tmp_stoi)
+        # stoi_sumpre.append(tmp_stoi)
+        # # breakpoint()
+
+        # 针对libritts采样率是24k的
+        min_len=min(rawwav.size()[1],prewav.size()[1])
+        rawwav_stoi=rawwav[:,:min_len].squeeze(0)
+        prewav_stoi=prewav[:,:min_len].squeeze(0)
+        tmp_stoi=stoi(rawwav_stoi.cpu(),prewav_stoi.cpu(),rawwav_sr,extended=False)
+        print("****stoi",tmp_stoi)
+        stoi_sumpre.append(tmp_stoi)
+
+    print("*************UTMOS_raw",utmos_sumgt,utmos_sumgt/len(preaudio))
+    print("*************UTMOS_encodec",utmos_sumgt,utmos_sumencodec/len(preaudio))
+    print("*************PESQ:",pesq_sumpre,pesq_sumpre/len(preaudio))
+    print("*************F1_score:",f1score_sumpre,f1score_sumpre/(len(preaudio)-f1score_filt),f1score_filt)
+    print("*************STOI:",np.mean(stoi_sumpre))
+    
+    
+
+if __name__=="__main__":
+    main()

metrics/periodicity.py

@@ -0,0 +1,105 @@
+import librosa
+import numpy as np
+import torch
+import torchaudio
+import torchcrepe
+from torchcrepe.loudness import REF_DB
+
+SILENCE_THRESHOLD = -60
+UNVOICED_THRESHOLD = 0.21
+
+"""
+Periodicity metrics adapted from https://github.com/descriptinc/cargan
+"""
+
+
+def predict_pitch(
+    audio: torch.Tensor, silence_threshold: float = SILENCE_THRESHOLD, unvoiced_treshold: float = UNVOICED_THRESHOLD
+):
+    """
+    Predicts pitch and periodicity for the given audio.
+
+    Args:
+        audio (Tensor): The audio waveform.
+        silence_threshold (float): The threshold for silence detection.
+        unvoiced_treshold (float): The threshold for unvoiced detection.
+
+    Returns:
+        pitch (ndarray): The predicted pitch.
+        periodicity (ndarray): The predicted periodicity.
+    """
+    # torchcrepe inference
+    pitch, periodicity = torchcrepe.predict(
+        audio,
+        fmin=50.0,
+        fmax=550,
+        sample_rate=torchcrepe.SAMPLE_RATE,
+        model="full",
+        return_periodicity=True,
+        device=audio.device,
+        pad=False,
+    )
+    pitch = pitch.cpu().numpy()
+    periodicity = periodicity.cpu().numpy()
+
+    # Calculate dB-scaled spectrogram and set low energy frames to unvoiced
+    hop_length = torchcrepe.SAMPLE_RATE // 100  # default CREPE
+    stft = torchaudio.functional.spectrogram(
+        audio,
+        window=torch.hann_window(torchcrepe.WINDOW_SIZE, device=audio.device),
+        n_fft=torchcrepe.WINDOW_SIZE,
+        hop_length=hop_length,
+        win_length=torchcrepe.WINDOW_SIZE,
+        power=2,
+        normalized=False,
+        pad=0,
+        center=False,
+    )
+
+    # Perceptual weighting
+    freqs = librosa.fft_frequencies(sr=torchcrepe.SAMPLE_RATE, n_fft=torchcrepe.WINDOW_SIZE)
+    perceptual_stft = librosa.perceptual_weighting(stft.cpu().numpy(), freqs) - REF_DB
+    silence = perceptual_stft.mean(axis=1) < silence_threshold
+
+    periodicity[silence] = 0
+    pitch[periodicity < unvoiced_treshold] = torchcrepe.UNVOICED
+
+    return pitch, periodicity
+
+
+def calculate_periodicity_metrics(y: torch.Tensor, y_hat: torch.Tensor):
+    """
+    Calculates periodicity metrics for the predicted and true audio data.
+
+    Args:
+        y (Tensor): The true audio data.
+        y_hat (Tensor): The predicted audio data.
+
+    Returns:
+        periodicity_loss (float): The periodicity loss.
+        pitch_loss (float): The pitch loss.
+        f1 (float): The F1 score for voiced/unvoiced classification
+    """
+    true_pitch, true_periodicity = predict_pitch(y)
+    pred_pitch, pred_periodicity = predict_pitch(y_hat)
+
+    true_voiced = ~np.isnan(true_pitch)
+    pred_voiced = ~np.isnan(pred_pitch)
+
+    periodicity_loss = np.sqrt(((pred_periodicity - true_periodicity) ** 2).mean(axis=1)).mean()
+
+    # Update pitch rmse
+    voiced = true_voiced & pred_voiced
+    difference_cents = 1200 * (np.log2(true_pitch[voiced]) - np.log2(pred_pitch[voiced]))
+    pitch_loss = np.sqrt((difference_cents ** 2).mean())
+
+    # voiced/unvoiced precision and recall
+    true_positives = (true_voiced & pred_voiced).sum()
+    false_positives = (~true_voiced & pred_voiced).sum()
+    false_negatives = (true_voiced & ~pred_voiced).sum()
+
+    precision = true_positives / (true_positives + false_positives)
+    recall = true_positives / (true_positives + false_negatives)
+    f1 = 2 * precision * recall / (precision + recall)
+
+    return periodicity_loss, pitch_loss, f1