Official PyTorch implementation of "RVAE-EM: Generative speech dereverberation based on recurrent variational auto-encoder and convolutive transfer function"
# Usage:
# python dnsmos_local.py -t c:\temp\DNSChallenge4_Blindset -o DNSCh4_Blind.csv -p
#

import argparse
import concurrent.futures
import glob
import os

import librosa
import numpy as np
import numpy.polynomial.polynomial as poly
import onnxruntime as ort
import pandas as pd
import soundfile as sf
from requests import session
from tqdm import tqdm
from other import mean_std
SAMPLING_RATE = 16000
INPUT_LENGTH = 9.01

class ComputeScore:
    def __init__(self, primary_model_path, p808_model_path) -> None:
        self.onnx_sess = ort.InferenceSession(primary_model_path)
        self.p808_onnx_sess = ort.InferenceSession(p808_model_path)
        
    def audio_melspec(self, audio, n_mels=120, frame_size=320, hop_length=160, sr=16000, to_db=True):
        mel_spec = librosa.feature.melspectrogram(y=audio, sr=sr, n_fft=frame_size+1, hop_length=hop_length, n_mels=n_mels)
        if to_db:
            mel_spec = (librosa.power_to_db(mel_spec, ref=np.max)+40)/40
        return mel_spec.T

    def get_polyfit_val(self, sig, bak, ovr, is_personalized_MOS):
        if is_personalized_MOS:
            p_ovr = np.poly1d([-0.00533021,  0.005101  ,  1.18058466, -0.11236046])
            p_sig = np.poly1d([-0.01019296,  0.02751166,  1.19576786, -0.24348726])
            p_bak = np.poly1d([-0.04976499,  0.44276479, -0.1644611 ,  0.96883132])
        else:
            p_ovr = np.poly1d([-0.06766283,  1.11546468,  0.04602535])
            p_sig = np.poly1d([-0.08397278,  1.22083953,  0.0052439 ])
            p_bak = np.poly1d([-0.13166888,  1.60915514, -0.39604546])

        sig_poly = p_sig(sig)
        bak_poly = p_bak(bak)
        ovr_poly = p_ovr(ovr)

        return sig_poly, bak_poly, ovr_poly

    def __call__(self, fpath, sampling_rate, is_personalized_MOS):
        aud, input_fs = sf.read(fpath)
        fs = sampling_rate
        if input_fs != fs:
            audio = librosa.resample(aud, input_fs, fs)
        else:
            audio = aud
        actual_audio_len = len(audio)
        len_samples = int(INPUT_LENGTH*fs)
        while len(audio) < len_samples:
            audio = np.append(audio, audio)
        
        num_hops = int(np.floor(len(audio)/fs) - INPUT_LENGTH)+1
        hop_len_samples = fs
        predicted_mos_sig_seg_raw = []
        predicted_mos_bak_seg_raw = []
        predicted_mos_ovr_seg_raw = []
        predicted_mos_sig_seg = []
        predicted_mos_bak_seg = []
        predicted_mos_ovr_seg = []
        predicted_p808_mos = []

        for idx in range(num_hops):
            audio_seg = audio[int(idx*hop_len_samples) : int((idx+INPUT_LENGTH)*hop_len_samples)]
            if len(audio_seg) < len_samples:
                continue

            input_features = np.array(audio_seg).astype('float32')[np.newaxis,:]
            p808_input_features = np.array(self.audio_melspec(audio=audio_seg[:-160])).astype('float32')[np.newaxis, :, :]
            oi = {'input_1': input_features}
            p808_oi = {'input_1': p808_input_features}
            p808_mos = self.p808_onnx_sess.run(None, p808_oi)[0][0][0]
            mos_sig_raw,mos_bak_raw,mos_ovr_raw = self.onnx_sess.run(None, oi)[0][0]
            mos_sig,mos_bak,mos_ovr = self.get_polyfit_val(mos_sig_raw,mos_bak_raw,mos_ovr_raw,is_personalized_MOS)
            predicted_mos_sig_seg_raw.append(mos_sig_raw)
            predicted_mos_bak_seg_raw.append(mos_bak_raw)
            predicted_mos_ovr_seg_raw.append(mos_ovr_raw)
            predicted_mos_sig_seg.append(mos_sig)
            predicted_mos_bak_seg.append(mos_bak)
            predicted_mos_ovr_seg.append(mos_ovr)
            predicted_p808_mos.append(p808_mos)

        clip_dict = {'filename': fpath, 'len_in_sec': actual_audio_len/fs, 'sr':fs}
        clip_dict['num_hops'] = num_hops
        clip_dict['OVRL_raw'] = np.mean(predicted_mos_ovr_seg_raw)
        clip_dict['SIG_raw'] = np.mean(predicted_mos_sig_seg_raw)
        clip_dict['BAK_raw'] = np.mean(predicted_mos_bak_seg_raw)
        clip_dict['OVRL'] = np.mean(predicted_mos_ovr_seg)
        clip_dict['SIG'] = np.mean(predicted_mos_sig_seg)
        clip_dict['BAK'] = np.mean(predicted_mos_bak_seg)
        clip_dict['P808_MOS'] = np.mean(predicted_p808_mos)
        return clip_dict

def main(args):
    models = glob.glob(os.path.join(args.testset_dir, "*"))
    audio_clips_list = []
    p808_model_path = os.path.join('DNSMOS', 'model_v8.onnx')

    if args.personalized_MOS:
        primary_model_path = os.path.join('pDNSMOS', 'sig_bak_ovr.onnx')
    else:
        primary_model_path = os.path.join('DNSMOS', 'sig_bak_ovr.onnx')

    compute_score = ComputeScore(primary_model_path, p808_model_path)

    rows = []
    clips = []
    clips = glob.glob(os.path.join(args.testset_dir, "*.wav"))
    is_personalized_eval = args.personalized_MOS
    desired_fs = SAMPLING_RATE
    for m in tqdm(models):
        max_recursion_depth = 10
        audio_path = os.path.join(args.testset_dir, m)
        audio_clips_list = glob.glob(os.path.join(audio_path, "*.wav"))
        while len(audio_clips_list) == 0 and max_recursion_depth > 0:
            audio_path = os.path.join(audio_path, "**")
            audio_clips_list = glob.glob(os.path.join(audio_path, "*.wav"))
            max_recursion_depth -= 1
        clips.extend(audio_clips_list)

    with concurrent.futures.ThreadPoolExecutor() as executor:
        future_to_url = {executor.submit(compute_score, clip, desired_fs, is_personalized_eval): clip for clip in clips}
        for future in tqdm(concurrent.futures.as_completed(future_to_url)):
            clip = future_to_url[future]
            try:
                data = future.result()
            except Exception as exc:
                print('%r generated an exception: %s' % (clip, exc))
            else:
                rows.append(data)            

    df = pd.DataFrame(rows)
    with open(os.path.join(args.testset_dir, "ave_metrics_mos.txt"), "w") as f:
        print(
            "OVRL: {:.2f} ± {:.2f}".format(*mean_std(df['OVRL'].to_numpy())), file=f 
        )
        print(
            "SIG: {:.2f} ± {:.2f}".format(*mean_std(df['SIG'].to_numpy())), file=f 
        )
        print(
            "BAK: {:.2f} ± {:.2f}".format(*mean_std(df['BAK'].to_numpy())), file=f 
        )
    
    
    
    if args.csv_path:
        csv_path = args.csv_path
        df.to_csv(csv_path)
    else:
        print(df.describe())

if __name__=="__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('-t', "--testset_dir", default='.', 
                        help='Path to the dir containing audio clips in .wav to be evaluated')
    parser.add_argument('-o', "--csv_path", default=None, help='Dir to the csv that saves the results')
    parser.add_argument('-p', "--personalized_MOS", action='store_true', 
                        help='Flag to indicate if personalized MOS score is needed or regular')
    
    args = parser.parse_args()

    main(args)

DNSMOS/dnsmos_local.py
Audio-WestlakeU/RVAE-EM
import os

import numpy as np

import scipy.stats
from scipy.signal import butter, sosfilt

import torch

from pesq import pesq
from pystoi import stoi


def si_sdr_components(s_hat, s, n):
    """
    """
    # s_target
    alpha_s = np.dot(s_hat, s) / np.linalg.norm(s)**2
    s_target = alpha_s * s

    # e_noise
    alpha_n = np.dot(s_hat, n) / np.linalg.norm(n)**2
    e_noise = alpha_n * n

    # e_art
    e_art = s_hat - s_target - e_noise
    
    return s_target, e_noise, e_art

def energy_ratios(s_hat, s, n):
    """
    """
    s_target, e_noise, e_art = si_sdr_components(s_hat, s, n)

    si_sdr = 10*np.log10(np.linalg.norm(s_target)**2 / np.linalg.norm(e_noise + e_art)**2)
    si_sir = 10*np.log10(np.linalg.norm(s_target)**2 / np.linalg.norm(e_noise)**2)
    si_sar = 10*np.log10(np.linalg.norm(s_target)**2 / np.linalg.norm(e_art)**2)

    return si_sdr, si_sir, si_sar

def mean_conf_int(data, confidence=0.95):
    a = 1.0 * np.array(data)
    n = len(a)
    m, se = np.mean(a), scipy.stats.sem(a)
    h = se * scipy.stats.t.ppf((1 + confidence) / 2., n-1)
    return m, h

class Method():
    def __init__(self, name, base_dir, metrics):
        self.name = name
        self.base_dir = base_dir
        self.metrics = {} 
        
        for i in range(len(metrics)):
            metric = metrics[i]
            value = []
            self.metrics[metric] = value 
            
    def append(self, matric, value):
        self.metrics[matric].append(value)

    def get_mean_ci(self, metric):
        return mean_conf_int(np.array(self.metrics[metric]))

def hp_filter(signal, cut_off=80, order=10, sr=16000):
    factor = cut_off /sr * 2
    sos = butter(order, factor, 'hp', output='sos')
    filtered = sosfilt(sos, signal)
    return filtered

def si_sdr(s, s_hat):
    alpha = np.dot(s_hat, s)/np.linalg.norm(s)**2   
    sdr = 10*np.log10(np.linalg.norm(alpha*s)**2/np.linalg.norm(
        alpha*s - s_hat)**2)
    return sdr

def snr_dB(s,n):
    s_power = 1/len(s)*np.sum(s**2)
    n_power = 1/len(n)*np.sum(n**2)
    snr_dB = 10*np.log10(s_power/n_power)
    return snr_dB

def pad_spec(Y):
    T = Y.size(3)
    if T%64 !=0:
        num_pad = 64-T%64
    else:
        num_pad = 0
    pad2d = torch.nn.ZeroPad2d((0, num_pad, 0,0))
    return pad2d(Y)


def ensure_dir(file_path):
    directory = file_path
    if not os.path.exists(directory):
        os.makedirs(directory)


def print_metrics(x, y, x_hat_list, labels, sr=16000):
    _si_sdr_mix = si_sdr(x, y)
    _pesq_mix = pesq(sr, x, y, 'wb')
    _estoi_mix = stoi(x, y, sr, extended=True)
    print(f'Mixture:  PESQ: {_pesq_mix:.2f}, ESTOI: {_estoi_mix:.2f}, SI-SDR: {_si_sdr_mix:.2f}')
    for i, x_hat in enumerate(x_hat_list):
        _si_sdr = si_sdr(x, x_hat)
        _pesq = pesq(sr, x, x_hat, 'wb')
        _estoi = stoi(x, x_hat, sr, extended=True)
        print(f'{labels[i]}: {_pesq:.2f}, ESTOI: {_estoi:.2f}, SI-SDR: {_si_sdr:.2f}')

def mean_std(data):
    # data = data[~np.isnan(data)]
    mean = np.mean(data)
    std = np.std(data)
    return mean, std

def print_mean_std(data, decimal=2):
    data = np.array(data)
    # data = data[~np.isnan(data)]
    mean = np.mean(data)
    std = np.std(data)
    if decimal == 2:
        string = f'{mean:.2f} ± {std:.2f}'
    elif decimal == 1:
        string = f'{mean:.1f} ± {std:.1f}'
    return string

DNSMOS/other.py
Audio-WestlakeU/RVAE-EM
import torch
import torchaudio
import numpy as np


def stft(audio, nfft, hop, win):
    """
    Func: STFT
    Params:
        nfft: nfft
        hop: hop size
        win: type of STFT analysis window, str, 'hann' or 'hamming'
    """
    if win.upper() == "HANN":
        window = torch.hann_window(nfft).to(audio.device)
    elif win.upper() == "HAMMING":
        window = torch.hamming_window(nfft).to(audio.device)

    spch_sptm = torch.stft(
        audio,
        n_fft=nfft,
        hop_length=hop,
        onesided=True,
        window=window,
        return_complex=True,
    )
    return spch_sptm


def istft(sptm, nfft, hop, win):
    """
    Func: iSTFT
    Params: same as stft
    """
    if win.upper() == "HANN":
        window = torch.hann_window(nfft).to(sptm.device)
    elif win.upper() == "HAMMING":
        window = torch.hamming_window(nfft).to(sptm.device)

    audio = torch.istft(sptm, n_fft=nfft, hop_length=hop, onesided=True, window=window)

    return audio.to("cpu")


def aud2sptm_woDC(audio, nfft, hop, win, *args, **kwargs):
    """
    Func: from audio to magnitude spectrogram without DC component
    Return: magnitude and phase spectrograms
    Params:
        audio: audio [bs,samples]
        others: same as stft
    """
    sptm = stft(
        audio,
        nfft=nfft,
        hop=hop,
        win=win,
    )
    sptm = sptm[:, 1:, :].permute(0, 2, 1)
    return sptm.abs(), torch.angle(sptm)  # [bs, T, F]


def sptm_woDC2aud(mag_woDC, phase_woDC, nfft, hop, win, *args, **kwargs):
    """
    Func: from magnitude spectrogram without DC component to audio
    Return: Audio
    Params:
        mag_woDC: magnitude spectrogram
        phase_woDC: phase spectrogram
        others: same as istft
    """
    mag_woDC = mag_woDC.permute(0, 2, 1)  # [bs,F,T]
    phase_woDC = phase_woDC.permute(0, 2, 1)
    bs, F, T = mag_woDC.shape
    sptm_woDC = mag_woDC * torch.exp(1j * phase_woDC)
    F = F + 1
    sptm = torch.zeros([bs, F, T]).to(sptm_woDC.device) + 0j
    sptm[:, 1:, :] = sptm_woDC
    wav = istft(sptm, nfft=nfft, hop=hop, win=win)
    return wav  # [bs,sample]


def audio2EMinput_woDC_preprocess(audio, nfft, hop, win, *args, **kwargs):
    """
    Func: from audio to RVAE-EM input
    Params:
        same as aud2sptm_woDC
    """
    sptm = stft(
        audio,
        nfft=nfft,
        hop=hop,
        win=win,
    )
    sptm_wo_DC = sptm[:, 1:, :]
    return sptm_wo_DC


def EMoutput_woDC2audio_postprocess(sptm_wo_DC, nfft, hop, win, *args, **kwargs):
    """
    Func: From EM output to audio
    Params:
        sptm_wo_DC: RVAE-EM output, complex spectrogram
        others: same as istft
    """
    bs, F, T = sptm_wo_DC.shape
    F = F + 1
    sptm_w_DC = torch.zeros([bs, F, T]).to(sptm_wo_DC.device) + 0j
    sptm_w_DC[:, 1:, :] = sptm_wo_DC

    wav = istft(sptm_w_DC, nfft=nfft, hop=hop, win=win)
    return wav

dataset/io.py
Audio-WestlakeU/RVAE-EM
import torch
from torch.utils.data import Dataset
import random
import scipy.signal
import os
import numpy as np
import soundfile as sf


class MyDataset(Dataset):
    def __init__(
        self,
        spch_index_txt: str,
        rir_index_txt: str,
        rir_target_index_txt: str,
        fs: int = 16000,
        spch_len: float = 5.104,
        norm_mode: str = "noisy_time",
        **kwargs
    ):
        """
        Class: dataset for training and validating
        Params:
            spch_index_txt: clean speech path saved as .txt
            rir_index_txt: reverberant RIR path saved as .txt
            rir_target_index_txt: dry RIR path savd as .txt
            fs: sample frequency
            spch_len: segement length (default 5.104s (320frames))
            norm_mode: normalization mode, noisy_time or time
        """
        self.fs = fs
        self.spch_len = spch_len

        assert (
            os.path.splitext(spch_index_txt)[-1] == ".txt"
            and os.path.splitext(rir_index_txt)[-1] == ".txt"
            and os.path.splitext(rir_target_index_txt)[-1] == ".txt"
        )

        # 读取clean speech列表
        spch_list = []
        spch_f = open(spch_index_txt, "r")
        for i in spch_f.readlines():
            spch_list.append(i.rstrip())

        # 读取rir列表
        rir_list = []
        rir_f = open(rir_index_txt, "r")
        for i in rir_f.readlines():
            rir_list.append(i.rstrip())

        rir_target_list = []
        rir_target_f = open(rir_target_index_txt, "r")
        for i in rir_target_f.readlines():
            rir_target_list.append(i.rstrip())

        # 实际使用的音频与RIR列表
        self.rir_list = rir_list
        self.rir_target_list = rir_target_list
        self.valid_seq_list = self._get_valid_speech_list(
            spch_list,
        )

        self.norm_mode = norm_mode

    def _get_valid_speech_list(
        self,
        spch_filename_list,
    ):
        valid_seq_list = []  # 可用语音List
        for wavfile in spch_filename_list:
            # 读取音频
            x, fs_x = sf.read(wavfile, dtype="float32")
            x = x.squeeze()
            # 判断采样频率是否正确
            assert self.fs == fs_x
            # 计算音频开始、结束位置
            idx_beg = 0
            idx_end = len(x)
            # 音频切分
            file_len = idx_end - idx_beg  # 非沉默音频段长度
            samples = int(self.fs * self.spch_len)
            # seq_len = (seq_len - 1) * hop
            n_seq = int(file_len // samples)
            for i in range(n_seq):
                seq_start = int(i * samples + idx_beg)
                seq_end = int((i + 1) * samples + idx_beg)
                seq_info = (wavfile, seq_start, seq_end)
                valid_seq_list.append(seq_info)

        return valid_seq_list

    def __len__(self):

        return len(self.valid_seq_list)

    def __getitem__(self, index):
        """
        读取音频
        """
        rir_index_list = torch.randint(
            0, len(self.rir_list), (len(self.valid_seq_list),)
        )

        rir_list = [self.rir_list[i] for i in rir_index_list]
        rir_target_list = [self.rir_target_list[i] for i in rir_index_list]

        ##################################################################
        # 读取clean speech
        spch_filename, seq_start, seq_end = self.valid_seq_list[index]
        spch_ori, fs_spch = sf.read(spch_filename, dtype="float32")
        assert self.fs == fs_spch
        spch_ori = spch_ori.squeeze()

        shift_max = min(seq_start - 0, spch_ori.shape[0] - seq_end, self.fs)
        shift = random.randint(-shift_max, shift_max)

        spch_ori = spch_ori[int(seq_start + shift) : int(seq_end + shift)]

        ##################################################################
        samples = int(self.fs * self.spch_len)

        rir, fs_rir = sf.read(rir_list[index], dtype="float32")
        assert self.fs == fs_rir
        rir = rir.squeeze()
        spch_noisy = scipy.signal.fftconvolve(spch_ori, rir, mode="full")
        spch_noisy = spch_noisy[:samples]

        rir_target, fs_rir_target = sf.read(rir_target_list[index], dtype="float32")
        assert self.fs == fs_rir_target
        rir_target = rir_target.squeeze()
        spch_dry = scipy.signal.fftconvolve(spch_ori, rir_target, mode="full")
        spch_dry = spch_dry[:samples]

        # normalization
        if np.max(np.abs(spch_noisy)) > 0:
            scale = np.max(np.abs(spch_noisy))
            spch_noisy = spch_noisy / scale
            spch_dry = spch_dry / scale
        if self.norm_mode.upper() == "TIME":
            spch_dry = spch_dry / np.max(np.abs(spch_dry))

        return torch.Tensor(spch_noisy), torch.Tensor(spch_dry)

dataset/mydataset.py
Audio-WestlakeU/RVAE-EM
from torch.utils.data import Dataset
import os
import numpy as np
import soundfile as sf
import glob


class TestDataset(Dataset):
    def __init__(self, spch_dir: str, fs: int = 16000, **kwargs):
        """
        Class: testset
        Params:
            spch_dir: Reverberant speech folder path
            fs: sample frequency
        """
        self.fs = fs
        assert os.path.exists(spch_dir)

        self.wav_files = glob.glob(os.path.join(spch_dir, "*.wav"))

    def __len__(self):

        return len(self.wav_files)

    def __getitem__(self, index):

        input, fs_spch = sf.read(self.wav_files[index], dtype="float32")
        assert self.fs == fs_spch

        filename = os.path.basename(self.wav_files[index])
        if np.max(np.abs(input)) > 0:
            spch_scale = np.max(np.abs(input))
            input = input / spch_scale * 0.95

        return input, filename


if __name__ == "__main__":

    pass

dataset/testdataset.py
Audio-WestlakeU/RVAE-EM
import os
import argparse
import json
import torch
import time
import warnings
import numpy as np
import random
from glob import glob
from math import ceil
import logging
from torchaudio import save
from torchinfo import summary
from torch.utils.data import DistributedSampler, DataLoader
import torch.multiprocessing as mp
from torch.distributed import init_process_group
from datetime import datetime
from utils.env import build_env, dict_to_namespace
from utils.utils import (
    load_checkpoint,
)
from model.RVAE import RVAE as MODEL
from dataset.testdataset import TestDataset as dataset
from dataset.io import audio2EMinput_woDC_preprocess as data_preprocess
from dataset.io import EMoutput_woDC2audio_postprocess as data_postprocess
from model.my_EM import MyEM

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
warnings.simplefilter(action="ignore", category=UserWarning)

logger = logging.getLogger("mylogger")
logger.setLevel(logging.DEBUG)


def inference(rank, a, c):

    # create namespace of args and configs
    a = dict_to_namespace(a)
    c = dict_to_namespace(c)

    # logger
    if rank == 0:
        fh = logging.FileHandler(os.path.join(a.save_path, "log.txt"), mode="a")
        fh.setLevel(logging.DEBUG)
        ch = logging.StreamHandler()
        ch.setLevel(logging.DEBUG)
        formatter = logging.Formatter("%(asctime)s %(levelname)s: %(message)s")
        fh.setFormatter(formatter)
        ch.setFormatter(formatter)
        logger.addHandler(fh)
        logger.addHandler(ch)

    # seed
    seed = c.seed
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    seed_this_gpu = c.seed + rank
    torch.cuda.manual_seed(seed_this_gpu)

    # create paths
    a.enhanced_path = a.save_path
    os.makedirs(a.enhanced_path, exist_ok=True)

    # DDP initialization
    init_process_group(
        backend=c.dist_config.dist_backend,
        init_method=c.dist_config.dist_url,
        world_size=c.dist_config.world_size * c.num_gpus,
        rank=rank,
    )

    # model initialization
    device = torch.device("cuda:{:d}".format(rank))
    torch.cuda.set_device(rank)
    torch.cuda.empty_cache()
    model = MODEL(**vars(c.model)).to(device)

    already_exist_wavs_filename = []
    for filepath_abs in glob(os.path.join(a.enhanced_path, "*")):
        if filepath_abs.endswith("wav"):
            already_exist_wavs_filename.append(os.path.basename(filepath_abs))

    logger.info("loading checkpoint from ", c.ckpt)
    state_dict = load_checkpoint(c.ckpt, device)
    model.load_state_dict(state_dict["params"])
    for name, param in model.named_parameters():
        param.requires_grad = False

    if rank == 0:
        logger.info(summary(model))

    # dataloader
    testset = dataset(**vars(c.data_setting_test), **vars(c.stft_setting))
    test_sampler = DistributedSampler(testset) if c.num_gpus > 1 else None
    test_loader = DataLoader(
        testset,
        num_workers=c.num_workers,
        shuffle=False,
        sampler=test_sampler,
        batch_size=c.batch_size[2],
        pin_memory=False,
        drop_last=False,
    )

    model.eval()
    algo = MyEM(model, vars(c.EM_kwargs))
    del model
    torch.cuda.empty_cache()
    start_time = time.time()

    # Inference
    with torch.cuda.amp.autocast():
        for i, batch in enumerate(test_loader):
            audio_input, filename = batch
            bs = audio_input.shape[0]
            if True:
                audio_input = audio_input.to(device)
                seq_len = audio_input.shape[1]
                seq_len_input = ceil(seq_len / c.stft_setting.hop) * c.stft_setting.hop
                audio_input_pad = torch.zeros([bs, seq_len_input]).to(device)
                audio_input_pad[:, :seq_len] = audio_input
                input = data_preprocess(audio_input_pad, **vars(c.stft_setting))
                output_sptm = algo.EM(input)
                output_wav = data_postprocess(output_sptm, **vars(c.stft_setting)).to(
                    "cpu"
                )
                output_wav = output_wav[:, :seq_len]
                for isample in range(bs):
                    save_path = os.path.join(a.enhanced_path, filename[isample])
                    save(
                        save_path,
                        (
                            output_wav[isample] / output_wav[isample].abs().max()
                        ).unsqueeze(0),
                        c.stft_setting.fs,
                    )
            if rank == 0:
                print(
                    "\r",
                    int(i + 1) / test_loader.__len__() * 100,
                    "%",
                    int((time.time() - start_time) / 60),
                    "min",
                    end="",
                    flush=True,
                )
    logger.info(
        "Total inference time: {:.1f} minutes".format((time.time() - start_time) / 60)
    )


def main():
    logger.info("Initializing Testing Process..")

    now = datetime.now()
    current_time = now.strftime("%Y-%m-%d-%H-%M")

    parser = argparse.ArgumentParser()

    parser.add_argument("--config", "-c", action="append", default=[])  # config
    parser.add_argument("--ckpt", required=True)  # checkpoint
    parser.add_argument(
        "--save_path", "-p", default="inferenced_model" + current_time
    )  # output folder

    args = parser.parse_args()

    json_config = {}
    for filename in args.config:
        with open(filename, "r") as f:  # 读取config
            json_config.update(json.load(f))

    config = dict_to_namespace(json_config)  # 转换为属性字典
    config.ckpt = args.ckpt
    build_env(config, "config.json", args.save_path)

    if torch.cuda.is_available():
        config.num_gpus = torch.cuda.device_count()
        config.batch_size[2] = int(config.batch_size[2] / config.num_gpus)
        logger.info("Batch size per GPU :", config.batch_size)
    else:
        raise ValueError("No GPU Devices!")

    if config.num_gpus > 1:
        mp.spawn(
            inference,
            nprocs=config.num_gpus,
            args=(
                vars(args),
                vars(config),
            ),
        )
    else:
        inference(0, vars(args), vars(config))


if __name__ == "__main__":
    main()

enhance.py
Audio-WestlakeU/RVAE-EM
from os.path import join
from glob import glob
from soundfile import read
from tqdm import tqdm
from pesq import pesq
import pandas as pd
import argparse
import torch
import numpy as np
import datetime
from pystoi import stoi
from utils.other import mean_std
from srmrpy import srmr
from torchmetrics.functional.audio import (
    scale_invariant_signal_distortion_ratio as sisdr,
)


def eval_metrics(a):

    output_root = a.out_path
    input_root = a.in_path
    ref_root = a.ref_path
    sr = a.sr
    device = "cpu"

    data = {
        "filename": [],
        "wbpesq": [],
        "nbpesq": [],
        "stoi": [],
        "estoi": [],
        "sisdr": [],
        "srmr": [],
    }

    # Evaluate standard metrics
    noisy_files = sorted(glob("{}/*.wav".format(input_root)))

    for noisy_file in tqdm(noisy_files, desc="Calculating metrics"):
        filename = noisy_file.split("/")[-1]
        x, sr_read = read(join(ref_root, filename))
        assert sr == sr_read
        x = x / np.max(x)
        x_method, _ = read(join(output_root, filename))
        x_method = x_method / np.max(x_method)

        data["filename"].append(filename)
        data["wbpesq"].append(pesq(sr, x, x_method, "wb"))
        data["nbpesq"].append(pesq(sr, x, x_method, "nb"))
        data["stoi"].append(stoi(x[: x_method.shape[-1]], x_method, sr, extended=False))
        data["estoi"].append(stoi(x[: x_method.shape[-1]], x_method, sr, extended=True))
        data["sisdr"].append(
            sisdr(torch.Tensor(x_method), torch.Tensor(x[: x_method.shape[-1]])).item()
        )
        data["srmr"].append(srmr(x_method, sr)[0])

    # Save results as DataFrame
    df = pd.DataFrame(data)

    # Print results
    with open(join(output_root, "ave_metrics.txt"), "w") as f:
        print(datetime.datetime.now(), file=f)
        print(output_root, file=f)
        print(
            "SISDR: {:.2f} ± {:.2f}".format(*mean_std(df["sisdr"].to_numpy())), file=f
        )
        print(
            "WBPESQ: {:.2f} ± {:.2f}".format(*mean_std(df["wbpesq"].to_numpy())), file=f
        )
        print(
            "NBPESQ: {:.2f} ± {:.2f}".format(*mean_std(df["nbpesq"].to_numpy())), file=f
        )
        print("STOI: {:.2f} ± {:.2f}".format(*mean_std(df["stoi"].to_numpy())), file=f)
        print(
            "ESTOI: {:.2f} ± {:.2f}".format(*mean_std(df["estoi"].to_numpy())), file=f
        )
        print("SRMR: {:.2f} ± {:.2f}".format(*mean_std(df["srmr"].to_numpy())), file=f)

    # Save DataFrame as csv file
    df.to_csv(join(output_root, "_results.csv"), index=False)


def main():
    print("Calculating the metrics of wavs..")

    parser = argparse.ArgumentParser()
    parser.add_argument("--out_path", "-o", required=True)  # output .wav folder
    parser.add_argument("--in_path", "-i", required=True)  # input .wav folder
    parser.add_argument("--ref_path", "-r", required=True)  # reference .wav folder
    parser.add_argument("--sr", default=16000)  # sample rate

    args = parser.parse_args()

    eval_metrics(args)


if __name__ == "__main__":
    main()

eval.py
Audio-WestlakeU/RVAE-EM
import torch
import torch.nn as nn
from .module import (
    build_conv2d,
    ResBlock2D,
    build_GRU,
    build_MLP,
    reparametrization,
)


class RVAE(nn.Module):
    def __init__(
        self,
        dim_x: int,
        dim_z: int,
        gru_dim_x_enc: int,
        gru_dim_z_enc: int,
        gru_dim_x_dec: int,
        pre_conv_enc: list,
        pre_conv_dec: list,
        resblock_enc: list,
        resblock_dec: list,
        post_conv_enc: list,
        post_conv_dec: list,
        num_resblock: int,
        num_GRU_layer_enc: int,
        num_GRU_layer_dec: int,
        dense_zmean_zlogvar: list,
        dense_activation_type: str,
        dropout_p: float,
        batch_norm: bool = False,
    ):
        super().__init__()
        self.net_pre_conv_enc = build_conv2d(pre_conv_enc, "Conv_pre_enc", False)
        self.net_conv_enc = ResBlock2D(
            resblock_enc, batch_norm, num_resblock, dropout_p
        )
        self.net_post_conv_enc = build_conv2d(post_conv_enc, "Conv_post_enc", False)
        self.gru_x_enc = build_GRU(dim_x, gru_dim_x_enc, num_GRU_layer_enc, True)
        self.gru_z_enc = build_GRU(dim_z, gru_dim_z_enc, num_GRU_layer_enc, False)
        dim_h_enc = gru_dim_x_enc * 2 + gru_dim_z_enc
        assert dense_zmean_zlogvar[-1] == dim_z
        self.mlp_zmean_enc, _ = build_MLP(
            dim_h_enc, dense_zmean_zlogvar, dense_activation_type, dropout_p
        )
        self.mlp_zlogvar_enc, _ = build_MLP(
            dim_h_enc, dense_zmean_zlogvar, dense_activation_type, dropout_p
        )
        self.gru_x_dec = build_GRU(dim_z, gru_dim_x_dec, num_GRU_layer_dec, True)
        self.net_pre_conv_dec = build_conv2d(pre_conv_dec, "Conv_pre_dec", False)
        self.net_conv_dec = ResBlock2D(
            resblock_dec, batch_norm, num_resblock, dropout_p
        )
        self.net_post_conv_dec = build_conv2d(post_conv_dec, "Conv_post_dec", False)
        self.num_GRU_layer_enc = num_GRU_layer_enc
        self.num_GRU_layer_dec = num_GRU_layer_dec
        self.dim_hz_enc = gru_dim_z_enc
        self.dim_z = dim_z

    def encoder(self, x):

        device = x.device
        x = (x + 1e-8).log()
        bs, seq_len, dim_feature = x.shape

        z = torch.zeros([bs, seq_len, self.dim_z]).to(device)
        z_t = torch.zeros([bs, self.dim_z]).to(device)
        h_hz_t_enc = torch.zeros([self.num_GRU_layer_enc, bs, self.dim_hz_enc]).to(
            device
        )

        zmean = torch.zeros([bs, seq_len, self.dim_z]).to(device)
        zlogvar = torch.zeros([bs, seq_len, self.dim_z]).to(device)

        # x的部分
        x = x.unsqueeze(1)
        x_temp = self.net_pre_conv_enc(x)
        x_temp = self.net_conv_enc(x_temp)
        x_temp = self.net_post_conv_enc(x_temp)
        hx_in_enc = x_temp.squeeze(1)
        hx_enc, _ = self.gru_x_enc(torch.flip(hx_in_enc, [1]))
        hx_enc = torch.flip(hx_enc, [1])

        # z和h_enc的部分
        for t in range(seq_len):
            hz_t_in_enc = z_t.unsqueeze(1)
            hz_t_enc, h_hz_t_enc = self.gru_z_enc(hz_t_in_enc, h_hz_t_enc)
            hz_t_enc = hz_t_enc.squeeze(1)
            h_t_enc = torch.cat([hx_enc[:, t, :], hz_t_enc], -1)
            zmean_t = self.mlp_zmean_enc(h_t_enc)
            zlogvar_t = self.mlp_zlogvar_enc(h_t_enc)
            z_t = reparametrization(zmean_t, zlogvar_t)
            z[:, t, :] = z_t
            zmean[:, t, :] = zmean_t
            zlogvar[:, t, :] = zlogvar_t

        return z, zmean, zlogvar

    def decoder(self, z):
        [bs, seq_len, _] = z.shape
        h_dec, _ = self.gru_x_dec(z)
        h_dec = h_dec.reshape([bs, seq_len, 2, -1])
        h_dec = h_dec.permute(0, 2, 1, 3)

        h_dec = self.net_pre_conv_dec(h_dec)
        h_dec = self.net_conv_dec(h_dec)
        h_dec = self.net_post_conv_dec(h_dec)
        logx = h_dec.squeeze(1)
        x = logx.exp()
        return x

    def forward(self, x):

        z, zmean, zlogvar = self.encoder(x)
        x_reconstruct = self.decoder(z)
        return x_reconstruct, zmean, zlogvar, z

model/RVAE.py
Audio-WestlakeU/RVAE-EM
import torch


class LossF:
    def loss_ISD(self, output: torch.Tensor, target: torch.Tensor, eps: float = 1e-8):
        """
        Return: IS loss
        Params:
            output: magnitude spectrogram (>=0), output of RVAE decoder, shape of [bs,T,F]
            target: magnitude spectrogram (>=0), clean, shape of [bs,T,F]
            eps: small constant, avoid numerical problems
        """
        bs, n_seq, _ = output.shape
        target = target**2 + eps
        output = output**2 + eps
        ret = torch.sum(target / output - torch.log(target / output) - 1)

        return ret / bs / n_seq

    def loss_KLD(self, zmean: torch.Tensor, zlogvar: torch.Tensor):
        """
        Return: KL loss
        Params:
            zmean: mean of latent variables, output of RVAE encoder, shape of [bs,T,D]
            zlogvar: log variance of latent variables, output of RVAE encoder, shape of [bs,T,D]
        """
        bs, n_seq, _ = zmean.shape
        zmean_p = torch.zeros_like(zmean)
        zlogvar_p = torch.zeros_like(zlogvar)
        ret = -0.5 * torch.sum(
            zlogvar
            - zlogvar_p
            - torch.div(zlogvar.exp() + (zmean - zmean_p).pow(2), zlogvar_p.exp())
            + 1
        )
        return ret / bs / n_seq

    def cal_KL_scale(
        self,
        cur_step: int,
        beta: float,
        zero_step: int,
        warmup_step: int,
        hold_step: int,
    ):
        """
        Return: the scale of KL loss
        Params:
            cur_step: current training step
            beta: base scale (default: 1)
            zero_step: keep scale = 0
            warmup_step: scale linear increased
            hold_step: keep scale=1
        """
        period = warmup_step + hold_step
        if cur_step < zero_step:
            return 0
        else:
            epoch_mod = (cur_step - zero_step) % period
            if epoch_mod < warmup_step:
                return beta * epoch_mod / warmup_step
            else:
                return beta

model/lossF.py
Audio-WestlakeU/RVAE-EM
import torch
import torch.nn as nn
import torch.nn.functional as F
from collections import OrderedDict


def init_weights(m, mean: float = 0.0, std: float = 0.01):
    """
    Initialize parameters in convolutional and transpose convolutional layers (zero-mean Gaussian)
    Params:
        mean: mean
        std: standard deviation
    """
    classname = m.__class__.__name__
    if classname.find("Conv") != -1:
        m.weight.data.normal_(mean, std)


def activation_func(type: str):
    """
    Return: activation function
    Params:
        type: 'relu', 'tanh', or 'leakyrelu'
    """
    if type.upper() == "RELU":
        nn_activation = nn.ReLU(inplace=False)
    elif type.upper() == "TANH":
        nn_activation = nn.Tanh()
    elif type.upper() == "LEAKYRELU":
        nn_activation = nn.LeakyReLU(0.1, inplace=False)
    else:
        raise ValueError("Unavailable activation type")
    return nn_activation


def build_MLP(dim_in: int, dense: list, activation_type: str, dropout_p: float = 0):
    """
    Return MLP and its feature dim of output
    Params:
        dim_in: input feature dim
        dense: list of layer dims
        activation_type: 'relu', 'tanh', or 'leakyrelu'
        dropout_p: dropout probability (not available for the output layer)
    """
    nn_activation = activation_func(activation_type)
    dic_layers = OrderedDict()
    if len(dense) == 0:
        dic_layers["Identity"] = nn.Identity()
        dim_y = dim_in
    else:
        for n in range(len(dense)):
            if n == 0:
                dic_layers["Linear" + str(n)] = nn.Linear(dim_in, dense[n])
            else:
                dic_layers["Linear" + str(n)] = nn.Linear(dense[n - 1], dense[n])
            if n != len(dense) - 1:
                dic_layers["activation" + str(n)] = nn_activation
                dic_layers["dropout" + str(n)] = nn.Dropout(p=dropout_p)
        dim_y = dense[-1]

    return nn.Sequential(dic_layers), dim_y


def build_conv2d(config: list, name: str, norm: bool = False):
    """
    Return: convolutional layer sequence
    Params:
        config: list of parameters, [input channel, output channel, kernel size, stride, padding]
        name: name of built CNN
        norm: use batch norm or not (default: no batch norm)
    """
    dic_layers = OrderedDict()
    [ch_in, ch_out, kernel_size, stride, padding] = config
    dic_layers[name] = nn.Conv2d(
        ch_in,
        ch_out,
        (kernel_size, kernel_size),
        (stride, stride),
        (padding, padding),
    )
    if norm:
        dic_layers[name + "_bn"] = torch.nn.BatchNorm2d(
            num_features=ch_out, affine=True, track_running_stats=True
        )
    return nn.Sequential(dic_layers)


class ResBlock2D(torch.nn.Module):
    """
    Resblock class in our work
    """

    def __init__(
        self, config: list, norm: bool = False, nblock: int = 1, dropout_p: float = 0
    ):
        """
        Params:
            config: list of parameters, [[input channel, output channel, kernel size, stride, padding]]
            norm: use batch norm or not (default: no batch norm)
            nblock: number of repeated blocks
            dropout_p: dropout probability (not available for the output layer)
        """
        super(ResBlock2D, self).__init__()
        self.c = config
        self.convs = nn.ModuleList([])
        self.n_layers = len(config)
        self.norm = norm
        for iblock in range(nblock):
            for n in range(self.n_layers):
                [ch_in, ch_out, kernel_size, stride, padding] = config[n]
                self.convs.append(
                    nn.Conv2d(
                        ch_in,
                        ch_out,
                        (kernel_size, kernel_size),
                        (stride, stride),
                        (padding, padding),
                    )
                )
        self.convs.apply(init_weights)
        if norm:
            self.bn = nn.ModuleList([])
            for iblock in range(nblock):
                [_, ch, _, _, _] = config[0]
                self.bn.append(
                    torch.nn.BatchNorm2d(
                        num_features=ch, affine=True, track_running_stats=True
                    )
                )
        self.dropout = nn.Dropout(dropout_p)

    def forward(self, x):
        i = 0
        x_res = x
        for conv in self.convs:
            iblock = i // self.n_layers  # 当前是第几个block
            i = i + 1
            xt = F.leaky_relu(x, 0.1, inplace=False)
            x = conv(xt)
            x = self.dropout(x)

            if i % self.n_layers == 0:
                if self.norm:  # use batch normalization
                    bn = self.bn[iblock]
                    x = bn(x)
                x = x + x_res
                x_res = x
        return x


def build_GRU(dim_in: int, dim_hidden: int, num_layers: int, bidir: bool):
    """
    Return: GRU module
    Params:
        dim_in: input layer dim
        dim_hidden: hidden layer dim
        num_layers: number of layers
        bidir: bidirectional GRU or forward GRU
    """
    gru = nn.GRU(
        input_size=dim_in,
        hidden_size=dim_hidden,
        num_layers=num_layers,
        batch_first=True,
        bidirectional=bidir,
    )
    gru.flatten_parameters()
    return gru


def reparametrization(mean, logvar):
    """
    Return: sampled latent variables
    Params:
        mean: mean of latent variables
        logvar: log variance of latent variables
    """
    std = torch.exp(0.5 * logvar)
    eps = torch.randn_like(std)
    return torch.addcmul(mean, eps, std)

model/module.py
Audio-WestlakeU/RVAE-EM
"""
实现自己算法中的采样与EM，实现去混响
"""
import torch
import torch.nn.functional
from torch import vmap


class MyEM:
    def __init__(self, model, EM_kparams, chunk_size=None):
        """
        Class: EM algorithm
        Params:
            model: RVAE network
            EM kparams: parameter dict
                CTF_len: CTF filter length
                max_steps: max iterations
                noi_pow_scale: initial Gaussian noise power
            chuck_size: chunk size for torch.vmap. Try smaller parameter when facing memory issues.
        """
        super().__init__()

        self.model = model
        del model
        self.noi_pow_scale = EM_kparams["noi_pow_scale"]
        self.CTF_P = EM_kparams["CTF_len"]
        self.max_steps = EM_kparams["max_steps"]
        self.update_speech_post_func = vmap(
            self.update_speech_post, chunk_size=chunk_size
        )
        self.update_CTF_func = vmap(self.update_CTF, chunk_size=chunk_size)
        self.update_noise_var_func = vmap(self.update_noise_var, chunk_size=chunk_size)

    def EM(self, Observed_data):
        """
        Func: EM algorithm
        """
        [self.bs, self.F, self.T] = Observed_data.shape
        self.device = Observed_data.device
        self.model = self.model.to(self.device)

        # Initialize the CTF filter and noise variance
        CTF, Noi_var = self.init_CTF_noi_var(
            Observed_data
        )  # [bs,F,P+1] complex,[bs,F] real
        Noi_var_reshape = Noi_var.reshape([-1])  # [bs*F], real
        Observed_data_reshape = Observed_data.reshape([-1, self.T])  # [bs*F,T], complex
        CTF_reshape = CTF.reshape([-1, self.CTF_P + 1])  # [bs*F,P+1], complex

        # 初始化隐变量
        Speech_prior_var, _, _, _ = self.model(Observed_data.abs().permute(0, 2, 1))
        Speech_prior_var = Speech_prior_var.permute(0, 2, 1).pow(2)
        Speech_prior_var_reshape = Speech_prior_var.reshape(
            [-1, self.T]
        )  # [bs*F,T], real

        for _ in range(self.max_steps):

            # E step: get the posterior of clean speech
            (
                Speech_post_mean_reshape,
                Speech_post_var_reshape,
            ) = self.update_speech_post_func(
                Observed_data_reshape,
                CTF_reshape,
                Speech_prior_var_reshape,
                Noi_var_reshape,
            )  # [bs*F,T] complex, [bs*F,T,T] real

            # 更新CTF参数
            CTF_reshape = self.update_CTF_func(
                Observed_data_reshape, Speech_post_mean_reshape, Speech_post_var_reshape
            )
            # 更新Sigma_noi参数
            Noi_var_reshape = self.update_noise_var_func(
                Observed_data_reshape,
                CTF_reshape,
                Speech_post_mean_reshape,
                Speech_post_var_reshape,
            )

        return Speech_post_mean_reshape.reshape([self.bs, self.F, self.T])

    def init_CTF_noi_var(self, Observed_data):
        """
        Func: Initialize CTF filter and noise variance
        Params:
            Observed_data: observation, shape of [bs, F, T]
        """
        bs, F, T = Observed_data.shape
        device = self.device

        Y_power = Observed_data.abs().pow(2).mean(2)  # [bs,F]
        Sigma_noi = Y_power * self.noi_pow_scale  # [bs,F]

        H_init = torch.zeros([bs, F, self.CTF_P + 1]) + 0j
        H_init[:, :, 0] = H_init[:, :, 0] + 1

        return H_init.to(device), Sigma_noi.to(device)  # [bs,F,P+1],[bs,F]

    def update_speech_post(self, Observed_data, CTF_filter, speech_var, noise_var):
        """
        Func: E-step, update posterior of clean speech
        Params:
            Observed_data: reverberant recordings, shape of [bs, F, T]
            CTF_filter: CTF filter, shape of [bs, F, P+1]
            speech_var: prior variance of clean speech
            noise_var: noise variance
        """

        device = self.device
        T = Observed_data.shape[0]
        P = self.CTF_P

        H_tilde = torch.zeros([T, T])  # [T,T]
        H_tilde = H_tilde + 0j
        H_tilde = H_tilde.to(device)

        speech_var = speech_var + 0j

        for p in range(P + 1):  # [T,T]
            H_tilde = H_tilde + CTF_filter[p] * torch.diag_embed(
                torch.ones_like(H_tilde.diag(-p)), -p
            )

        UBVA = noise_var.pow(-1) * torch.mm(
            torch.mm(H_tilde.H, H_tilde), speech_var.diag_embed(0)
        )

        Sigma_f = speech_var.diag_embed(0) - torch.mm(
            torch.mm(speech_var.diag_embed(0), UBVA),
            ((torch.eye(UBVA.shape[0], device=UBVA.device) + UBVA).inverse()),
        )

        Mu_f = noise_var.pow(-1) * torch.mm(
            torch.mm(Sigma_f, H_tilde.H), Observed_data.unsqueeze(1)
        ).squeeze(
            1
        )  # [T]

        return Mu_f, Sigma_f  # [T] complex, [T,T] complex

    def update_CTF(self, Observed_data, Post_mean, Post_var):
        """
        Func: M-step, update CTF filter, shape of [P+1], complex
        Params:
            Observed_data: observation in single band [T], complex
            Post_mean: posterior mean of clean speech in single band [T], complex
            Post_var: posterior variance of clean speech in single band [T,T], complex
        """
        numerator = (
            torch.zeros([1, self.CTF_P + 1]).to(self.device) + 0j
        )  # [1,P+1], complex
        denominator = (
            torch.zeros([self.CTF_P + 1, self.CTF_P + 1]).to(self.device) + 0j
        )  # [P+1,P+1], complex
        for t in range(self.CTF_P, self.T):
            Mu_t = (
                Post_mean[t - self.CTF_P : t + 1].flip(0).unsqueeze(0)
            )  # [1,P+1], complex
            Sigma_t = (
                Post_var[t - self.CTF_P : t + 1, t - self.CTF_P : t + 1].flip(0).flip(1)
            )  # [P+1,P+1], complex

            numerator = numerator + Observed_data[t] * Mu_t.conj()
            denominator = denominator + torch.mm(Mu_t.T, Mu_t.conj()) + Sigma_t

        H = torch.mm(numerator, denominator.inverse()).squeeze()  # [P+1], complex

        return H

    def update_noise_var(self, Observed_data, CTF_filter, Post_mean, Post_var):
        """
        Func: M-step, update noise variance in a single band
        Params:
            Observed_data: observation in a single band, [T] complex
            CTF_filter: CTF filter, [P+1] complex
            Post_mean: posterior mean of clean speech, [T] complex
            Post_var: posterior variance of clean speech, [T,T] complex
        """
        T = Observed_data.shape[0]
        H_tilde = torch.zeros([T, T]) + 0j  # [T,T] complex
        H_tilde = H_tilde.to(self.device)

        for p in range(self.CTF_P + 1):
            H_tilde = H_tilde + CTF_filter[p] * torch.diag_embed(
                torch.ones_like(H_tilde.diag(-p)), -p
            )

        numerator = torch.norm(
            Observed_data.unsqueeze(1) - torch.mm(H_tilde, Post_mean.unsqueeze(1)),
            "fro",
        ).pow(2)

        numerator = (
            numerator
            + torch.trace(torch.mm(torch.mm(H_tilde, Post_var), H_tilde.H)).real
        )

        return numerator / T + 1e-8

model/my_EM.py
Audio-WestlakeU/RVAE-EM
import os
import random
import gpuRIR
import numpy as np
import soundfile as sf
import argparse
from tqdm import tqdm
import shutil
from utils import read_yaml


def gen_RIR(room_sz, T60, pos_src, pos_rcv, fs):

    beta = gpuRIR.beta_SabineEstimation(room_sz, T60)
    num_images = gpuRIR.t2n(2, room_sz)
    RIR = gpuRIR.simulateRIR(
        room_sz, beta, pos_src, pos_rcv, num_images, 2, fs
    ).squeeze()

    RIR_target = gpuRIR.simulateRIR(
        room_sz, [0.01 for i in range(6)], pos_src, pos_rcv, num_images, 2, fs
    ).squeeze()

    return RIR, RIR_target


if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="Gen RIRs")

    parser.add_argument("-c", "--config", type=str, default="config/data/gen_data.yaml")
    args = parser.parse_args()

    config = read_yaml(args.config)

    RIR_save_root = config["RIR_save_root"]
    T60_range = config["T60_range"]
    room_sz_range = config["room_sz_range"]
    min_distance_to_wall = config["min_distance_to_wall"]
    num_RIRs = config["num_RIRs"]
    fs = config["sample_rate"]

    if config["random_seed"]:
        print("random seed: ", config["random_seed"])
        random.seed(config["random_seed"])

    RIR_save_root_noisy = os.path.join(RIR_save_root, "noisy")
    RIR_save_root_target = os.path.join(RIR_save_root, "target")

    os.makedirs(RIR_save_root, exist_ok=True)
    shutil.rmtree(RIR_save_root)
    os.makedirs(RIR_save_root_noisy, exist_ok=True)
    os.makedirs(RIR_save_root_target, exist_ok=True)

    RIRs_noisy = np.zeros((num_RIRs, 2 * fs))
    RIRs_target = np.zeros((num_RIRs, 2 * fs))

    T60s_noisy = np.zeros(num_RIRs)
    for i in tqdm(range(num_RIRs), desc="Generating"):

        room_sz = np.array(
            [
                np.random.uniform(
                    room_sz_range[0],
                    room_sz_range[1],
                ),
                np.random.uniform(
                    room_sz_range[2],
                    room_sz_range[3],
                ),
                np.random.uniform(
                    room_sz_range[4],
                    room_sz_range[5],
                ),
            ],
        )
        pos_src = np.array(
            [
                [
                    np.random.uniform(
                        min_distance_to_wall, room_sz[0] - min_distance_to_wall
                    ),
                    np.random.uniform(
                        min_distance_to_wall, room_sz[1] - min_distance_to_wall
                    ),
                    np.random.uniform(
                        min_distance_to_wall, room_sz[2] - min_distance_to_wall
                    ),
                ]
            ]
        )

        pos_rcv = np.array(
            [
                [
                    np.random.uniform(
                        min_distance_to_wall, room_sz[0] - min_distance_to_wall
                    ),
                    np.random.uniform(
                        min_distance_to_wall, room_sz[1] - min_distance_to_wall
                    ),
                    np.random.uniform(
                        min_distance_to_wall, room_sz[2] - min_distance_to_wall
                    ),
                ]
            ]
        )
        T60 = np.random.uniform(T60_range[0], T60_range[1])
        T60s_noisy[i] = T60
        RIRs_noisy[i], RIRs_target[i] = gen_RIR(room_sz, T60, pos_src, pos_rcv, fs)

        filename = "{0}_T60_{1:.2f}.wav".format(i, T60s_noisy[i])
        savepath_noisy = os.path.abspath(os.path.join(RIR_save_root_noisy, filename))
        savepath_target = os.path.abspath(os.path.join(RIR_save_root_target, filename))
        sf.write(savepath_noisy, RIRs_noisy[i], fs)
        sf.write(savepath_target, RIRs_target[i], fs)

prepare_data/gen_rirs.py
Audio-WestlakeU/RVAE-EM
import torch
from torch.utils.data import Dataset
import random
import scipy.signal
import os
import numpy as np
import soundfile as sf
from utils import read_yaml, wav_path2txt
import argparse


class MyDataset(Dataset):
    def __init__(
        self,
        spch_index_txt: str,
        rir_index_txt: str,
        rir_target_index_txt: str,
        fs: int = 16000,
        spch_len: float = 5.104,
        **kwargs
    ):
        """
        spch_path: 存放干净语音文件名的txt文件地址
        rir_path: 存放RIR文件名的txt文件地址
        fs: 采样率Hz
        spch_len: 音频持续时间s
        norm_mode: 归一化方式
        """
        self.fs = fs
        self.spch_len = spch_len

        assert (
            os.path.splitext(spch_index_txt)[-1] == ".txt"
            and os.path.splitext(rir_index_txt)[-1] == ".txt"
            and os.path.splitext(rir_target_index_txt)[-1] == ".txt"
        )

        # 读取clean speech列表
        spch_list = []
        spch_f = open(spch_index_txt, "r")
        for i in spch_f.readlines():
            spch_list.append(i.rstrip())

        # 读取rir列表
        rir_list = []
        rir_f = open(rir_index_txt, "r")
        for i in rir_f.readlines():
            rir_list.append(i.rstrip())

        rir_target_list = []
        rir_target_f = open(rir_target_index_txt, "r")
        for i in rir_target_f.readlines():
            rir_target_list.append(i.rstrip())

        # 实际使用的音频与RIR列表
        self.rir_list = rir_list
        self.rir_target_list = rir_target_list
        self.valid_seq_list = self._get_valid_speech_list(
            spch_list,
        )

    def _get_valid_speech_list(
        self,
        spch_filename_list,
    ):
        valid_seq_list = []  # 可用语音List
        for wavfile in spch_filename_list:
            # 读取音频
            x, fs_x = sf.read(wavfile, dtype="float32")
            x = x.squeeze()
            # 判断采样频率是否正确
            assert self.fs == fs_x
            # 计算音频开始、结束位置
            idx_beg = 0
            idx_end = len(x)
            # 音频切分
            file_len = idx_end - idx_beg  # 非沉默音频段长度
            samples = int(self.fs * self.spch_len)
            # seq_len = (seq_len - 1) * hop
            n_seq = int(file_len // samples)
            for i in range(n_seq):
                seq_start = int(i * samples + idx_beg)
                seq_end = int((i + 1) * samples + idx_beg)
                seq_info = (wavfile, seq_start, seq_end)
                valid_seq_list.append(seq_info)

        return valid_seq_list

    def __len__(self):

        return len(self.valid_seq_list)

    def __getitem__(self, index):
        """
        读取音频
        """
        rir_index_list = torch.randint(
            0, len(self.rir_list), (len(self.valid_seq_list),)
        )

        rir_list = [self.rir_list[i] for i in rir_index_list]
        rir_target_list = [self.rir_target_list[i] for i in rir_index_list]

        ##################################################################
        # 读取clean speech
        spch_filename, seq_start, seq_end = self.valid_seq_list[index]
        spch_ori, fs_spch = sf.read(spch_filename, dtype="float32")
        # 检查采样率
        assert self.fs == fs_spch
        spch_ori = spch_ori.squeeze()

        # shift_max = min(seq_start - 0, spch_ori.shape[0] - seq_end, self.fs)
        shift = 0

        spch_ori = spch_ori[int(seq_start + shift) : int(seq_end + shift)]

        ##################################################################
        samples = int(self.fs * self.spch_len)
        # 加入混响
        # 读取RIR
        rir, fs_rir = sf.read(rir_list[index], dtype="float32")
        assert self.fs == fs_rir
        rir = rir.squeeze()
        spch_noisy = scipy.signal.fftconvolve(spch_ori, rir, mode="full")
        spch_noisy = spch_noisy[:samples]

        rir_target, fs_rir_target = sf.read(rir_target_list[index], dtype="float32")
        assert self.fs == fs_rir_target
        rir_target = rir_target.squeeze()
        spch_dry = scipy.signal.fftconvolve(spch_ori, rir_target, mode="full")
        spch_dry = spch_dry[:samples]

        # 归一化
        if np.max(np.abs(spch_noisy)) > 0:
            scale = np.max(np.abs(spch_noisy))
            spch_noisy = spch_noisy / scale
            spch_dry = spch_dry / scale

        return spch_noisy, spch_dry


if __name__ == "__main__":
    """
    Gen datasets
    """
    from tqdm import tqdm

    parser = argparse.ArgumentParser(description="Gen clean and reverberant speech")

    parser.add_argument("-c", "--config", type=str, required=True)
    args = parser.parse_args()

    config = read_yaml(args.config)

    wav_clean_root = config["wav_clean_root"]
    RIR_save_root = config["RIR_save_root"]
    wav_save_root = config["wav_save_root"]
    fs = config["sample_rate"]
    spch_len = config["spch_len"]
    RIR_noisy_root = os.path.join(RIR_save_root, "noisy")
    RIR_target_root = os.path.join(RIR_save_root, "target")
    input_path = os.path.join(wav_save_root, "input")
    target_path = os.path.join(wav_save_root, "target")
    wav_path2txt(wav_clean_root)
    wav_path2txt(RIR_noisy_root)
    wav_path2txt(RIR_target_root)

    os.makedirs(wav_save_root, exist_ok=True)
    os.makedirs(input_path, exist_ok=True)
    os.makedirs(target_path, exist_ok=True)

    # SEED
    seed = config["random_seed"]
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.cuda.manual_seed(seed)

    dataset = MyDataset(
        spch_index_txt=os.path.join(wav_clean_root, "index.txt"),
        rir_index_txt=os.path.join(RIR_noisy_root, "index.txt"),
        rir_target_index_txt=os.path.join(RIR_target_root, "index.txt"),
        fs=fs,
        spch_len=spch_len,
    )

    for i, batch in tqdm(enumerate(dataset), desc="Generating"):
        audio_input, audio_target = batch
        filename = "sample{:06d}.wav".format(i)
        input_file_path = os.path.join(input_path, filename)
        target_file_path = os.path.join(target_path, filename)
        sf.write(input_file_path, audio_input, fs)
        sf.write(target_file_path, audio_target, fs)

prepare_data/gen_testset.py
Audio-WestlakeU/RVAE-EM
import yaml
from glob import glob
import os


def read_yaml(config_path):
    # 读取yaml文件
    with open(config_path, encoding="utf-8") as f:
        config = yaml.load(f.read(), Loader=yaml.FullLoader)
    return config


def wav_path2txt(root):
    # 获取wav文件的地址存入txt文件
    files = sorted(glob("{}/**/*.wav".format(root), recursive=True))
    f = open(os.path.join(root, "index.txt"), "w")
    for filename in files:
        f.write(filename + os.linesep)
    f.close()
    return None

prepare_data/utils.py
Audio-WestlakeU/RVAE-EM
import os
import sys
import argparse
import json
import warnings
import random
import numpy as np
import torch
import torch.multiprocessing as mp
import logging
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DistributedSampler, DataLoader, RandomSampler
from torch.distributed import init_process_group
from torch.nn.parallel import DistributedDataParallel
from torchinfo import summary
import time
from datetime import datetime
from pesq import pesq
from pystoi import stoi

from utils.env import build_env, dict_to_namespace
from utils.utils import (
    scan_checkpoint,
    load_checkpoint,
    set_optimizer,
    set_scheduler,
    save_checkpoint,
    plot_spectrogram,
    wav_path2txt,
)
from model.RVAE import RVAE as MODEL
from model.lossF import LossF
from dataset.mydataset import MyDataset as dataset
from dataset.io import aud2sptm_woDC as data_preprocess
from dataset.io import sptm_woDC2aud as data_postprocess

# global settings
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
warnings.simplefilter(action="ignore", category=UserWarning)

# CUDNN
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True

logger = logging.getLogger("mylogger")
logger.setLevel(logging.DEBUG)

# training
def train(rank, a, c):

    # create namespace of args and configs
    a = dict_to_namespace(a)
    c = dict_to_namespace(c)
    
    # logger
    if rank == 0:
        fh = logging.FileHandler(os.path.join(a.save_path, "log.txt"), mode="a")
        fh.setLevel(logging.DEBUG)
        ch = logging.StreamHandler()
        ch.setLevel(logging.DEBUG)
        formatter = logging.Formatter("%(asctime)s %(levelname)s: %(message)s")
        fh.setFormatter(formatter)
        ch.setFormatter(formatter)
        logger.addHandler(fh)
        logger.addHandler(ch)
        
    # seed
    seed = c.seed
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    seed_this_gpu = c.seed + rank
    torch.cuda.manual_seed(seed_this_gpu)

    # create paths
    os.makedirs(a.save_path, exist_ok=True)
    a.ckpt_save_path = os.path.join(a.save_path, "ckpt")
    os.makedirs(a.ckpt_save_path, exist_ok=True)

    # DDP initialization
    init_process_group(
        backend=c.dist_config.dist_backend,
        init_method=c.dist_config.dist_url,
        world_size=c.dist_config.world_size * c.num_gpus,
        rank=rank,
    )
    
        # model and optimizer initialization
    device = torch.device("cuda:{:d}".format(rank))
    torch.cuda.set_device(rank)
    torch.cuda.empty_cache()
    model = MODEL(**vars(c.model)).to(device)  # to device
    optim = set_optimizer(model.parameters(), c.optimizer)
    scheduler = set_scheduler(optim, c.lr_scheduler)
    state_dict = None
    steps = 0
    last_epoch = -1
    ckpt = scan_checkpoint(a.ckpt_save_path, "ckpt_")
    if a.start_ckpt:
        if rank == 0:
            logger.info("Loading pretrained model from {}".format(a.start_ckpt))
        state_dict = load_checkpoint(a.start_ckpt, device)
        model.load_state_dict(state_dict["params"])
    else:
        if ckpt:
            if rank == 0:
                logger.info(
                    "Resume training process. Loading checkpoint from {}".format(ckpt)
                )
            state_dict = load_checkpoint(ckpt, device)
            model.load_state_dict(state_dict["params"])
            steps = state_dict["steps"] + 1
            last_epoch = state_dict["epoch"]
            # optim.load_state_dict(state_dict["optim"])
            # scheduler.load_state_dict(state_dict["scheduler"])
        else:
            if rank == 0:
                logger.info("New training process. Saved to {}".format(a.save_path))
                
    if rank == 0:
        logger.info(summary(model))
    if c.num_gpus > 1:  # DDP model
        model = DistributedDataParallel(
            model,
            device_ids=[rank],
            find_unused_parameters=False,
            broadcast_buffers=False,
        )

    # dataloader
    trainset = dataset(**vars(c.data_setting_train), **vars(c.stft_setting))
    train_sampler = (
        DistributedSampler(trainset) if c.num_gpus > 1 else RandomSampler(trainset)
    )
    train_loader = DataLoader(
        trainset,
        num_workers=c.num_workers,
        shuffle=False,
        sampler=train_sampler,
        batch_size=c.batch_size[0],
        pin_memory=True,
        drop_last=True,
        prefetch_factor=4,
        persistent_workers=True,
    )
    if rank == 0:
        validset = dataset(**vars(c.data_setting_val), **vars(c.stft_setting))
        validation_loader = DataLoader(
            validset,
            num_workers=c.num_workers,
            shuffle=False,
            sampler=None,
            batch_size=c.batch_size[1],
            pin_memory=True,
            drop_last=True,
            prefetch_factor=4,
            persistent_workers=True,
        )
        sw = SummaryWriter(os.path.join(a.save_path, "logs"))
        
    # training process
    model.train()
    if c.use_amp:
        scaler = torch.cuda.amp.GradScaler()
    for epoch in range(max(0, last_epoch), a.training_epochs):
        if rank == 0:
            start = time.time()
            logger.info("Epoch: {}".format(epoch + 1))
        if c.num_gpus > 1:
            train_sampler.set_epoch(epoch)
        for i, batch in enumerate(train_loader):
            if rank == 0:
                start_b = time.time()
            optim.zero_grad()
            audio_input, audio_clean = batch
            audio_input = audio_input.to(device)
            audio_clean = audio_clean.to(device)
            input, input_phase = data_preprocess(audio_input, **vars(c.stft_setting))
            target, _ = data_preprocess(audio_clean, **vars(c.stft_setting))
            if c.use_amp:
                with torch.cuda.amp.autocast():
                    output, z_mean, z_logvar, _ = model(input)
                    loss_IS = LossF().loss_ISD(output, target)
                    loss_KL = LossF().loss_KLD(z_mean, z_logvar)
                    scale_KL = LossF().cal_KL_scale(
                        steps,
                        c.beta,
                        c.beta_zero_step,
                        c.beta_warmup_step,
                        c.beta_holdon_step,
                    )
                    loss_tot = loss_IS + loss_KL * scale_KL
                    assert torch.isnan(loss_tot) == False
                    scaler.scale(loss_tot).backward()
                    if c.gradient_clip:
                        torch.nn.utils.clip_grad.clip_grad_norm(
                            model.parameters(), max_norm=c.gradient_clip
                        )
                    scaler.step(optim)
                    scaler.update()
            else:
                output, z_mean, z_logvar, _ = model(input)
                loss_IS = LossF().loss_ISD(output, target)
                loss_KL = LossF().loss_KLD(z_mean, z_logvar)
                scale_KL = LossF().cal_KL_scale(
                    steps,
                    c.beta,
                    c.beta_zero_step,
                    c.beta_warmup_step,
                    c.beta_holdon_step,
                )
                loss_tot = loss_IS + loss_KL * scale_KL
                assert torch.isnan(loss_tot) == False
                loss_tot.backward()
                if c.gradient_clip:
                    torch.nn.utils.clip_grad.clip_grad_norm(
                        model.parameters(), max_norm=c.gradient_clip
                    )
                optim.step()
            if rank == 0:
                # stdout logger
                if steps % a.stdout_interval == 0:
                    logger.info(
                        "Steps:{:d}, Loss:{:.1f}, Loss_IS:{:.1f}, Loss_KL:{:.1f}, KL_scale:{:.1f}, lr:{:.2e}, time_cost:{:.1f}".format(
                            steps,
                            loss_tot,
                            loss_IS,
                            loss_KL,
                            scale_KL,
                            optim.param_groups[0]["lr"],
                            time.time() - start_b,
                        )
                    )

                # checkpointing
                if steps % a.checkpoint_interval == 0 and steps != 0:
                    checkpoint_filename = "ckpt_{:010d}".format(steps)
                    checkpoint_path = os.path.join(
                        a.ckpt_save_path, checkpoint_filename
                    )
                    save_checkpoint(
                        checkpoint_path,
                        {
                            "params": (
                                model.module if c.num_gpus > 1 else model
                            ).state_dict(),
                            "optim": optim.state_dict(),
                            "scheduler": scheduler.state_dict() if scheduler else None,
                            "steps": steps,
                            "epoch": epoch,
                        },
                    )

                # tensorboard summary logging
                if steps % a.summary_interval == 0:
                    sw.add_scalar("training/loss", loss_tot, steps)
                    sw.add_scalar("training/loss_IS", loss_IS, steps)
                    sw.add_scalar("training/loss_KL", loss_KL, steps)
                    sw.add_scalar("epoch", epoch, steps)
                    sw.add_scalar("lr", optim.param_groups[0]["lr"], steps)

                # validation
                if steps % a.validation_interval == 0:  # and steps != 0:
                    logger.info("validating ...")
                    model.eval()
                    torch.cuda.empty_cache()
                    val_loss_tot_tot = 0
                    val_loss_IS_tot = 0
                    val_loss_KL_tot = 0
                    metric_wbpesq_tot = 0
                    metric_estoi_tot = 0
                    tot_bs = 0
                    with torch.no_grad():
                        for j, batch in enumerate(validation_loader):
                            audio_input, audio_clean = batch
                            audio_input = audio_input.to(device)
                            audio_clean = audio_clean.to(device)
                            input, input_phase = data_preprocess(
                                audio_input, **vars(c.stft_setting)
                            )
                            target, _ = data_preprocess(
                                audio_clean, **vars(c.stft_setting)
                            )
                            bs = input.shape[0]

                            output, z_mean, z_logvar, _ = model(input)
                            loss_IS = LossF().loss_ISD(output, target)
                            loss_KL = LossF().loss_KLD(z_mean, z_logvar)
                            loss_tot = loss_IS + loss_KL

                            output_wavs = data_postprocess(
                                output, input_phase, **vars(c.stft_setting)
                            )

                            val_loss_tot_tot += loss_tot.item() * bs
                            val_loss_IS_tot += loss_IS.item() * bs
                            val_loss_KL_tot += loss_KL.item() * bs

                            tot_bs += bs

                            if j == 0 and rank == 0:
                                metric_wbpesq_tot = metric_wbpesq_tot + pesq(
                                    c.stft_setting.fs,
                                    audio_clean[0].cpu().numpy(),
                                    output_wavs[0].cpu().numpy(),
                                    "wb",
                                )
                                metric_estoi_tot = metric_estoi_tot + stoi(
                                    audio_clean[0].cpu().numpy(),
                                    output_wavs[0].cpu().numpy(),
                                    c.stft_setting.fs,
                                    extended=True,
                                )
                                spectrogram = (
                                    output[0].permute(1, 0).log().cpu().numpy()
                                )
                                sw.add_figure(
                                    "generated/y_hat_spec_{}".format(steps),
                                    plot_spectrogram(spectrogram),
                                    steps,
                                )
                                if steps == 0:
                                    spectrogram = (
                                        input[0].permute(1, 0).log().cpu().numpy()
                                    )
                                    sw.add_figure(
                                        "clean/y_spec_{}".format(steps),
                                        plot_spectrogram(spectrogram),
                                        steps,
                                    )

                        val_loss_tot = val_loss_tot_tot / tot_bs
                        val_loss_IS = val_loss_IS_tot / tot_bs
                        val_loss_KL = val_loss_KL_tot / tot_bs
                        metric_wbpesq = metric_wbpesq_tot
                        metric_estoi = metric_estoi_tot

                        sw.add_scalar("validation/loss", val_loss_tot, steps)
                        sw.add_scalar("validation/loss_IS", val_loss_IS, steps)
                        sw.add_scalar("validation/loss_KL", val_loss_KL, steps)
                        sw.add_scalar(
                            "validation/metric_wbpesq_1sample", metric_wbpesq, steps
                        )
                        sw.add_scalar(
                            "validation/metric_estoi_1sample", metric_estoi, steps
                        )

                    model.train()

            steps += 1
            
        if scheduler:
            scheduler.step()

        if rank == 0:
            print(
                "Time taken for epoch {} is {} sec\n".format(
                    epoch + 1, int(time.time() - start)
                )
            )

def main():

    now = datetime.now()
    current_time = now.strftime("%Y-%m-%d-%H-%M")

    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--config", "-c", action="append", default=[], help="config files (.json)"
    )
    parser.add_argument(
        "--save_path", "-p", default="saved_model" + current_time, help="saved path"
    )
    parser.add_argument(
        "--training_epochs", default=100000, type=int, help="maximum epochs"
    )
    parser.add_argument(
        "--stdout_interval", default=10, type=int, help="stdout interval (steps)"
    )
    parser.add_argument(
        "--checkpoint_interval", default=1000, type=int, help="save interval (steps)"
    )
    parser.add_argument(
        "--summary_interval",
        default=100,
        type=int,
        help="tensorboard summary interval (steps)",
    )
    parser.add_argument(
        "--validation_interval",
        default=1000,
        type=int,
        help="validation interval (steps)",
    )
    parser.add_argument(
        "--start_ckpt", "-r", type=str, required=False, help="pretrained model (.ckpt)"
    )

    args = parser.parse_args()

    json_config = {}
    for filename in args.config:
        with open(filename, "r") as f:  # 读取config
            json_config.update(json.load(f))

    config = dict_to_namespace(json_config)  # 转换为属性字典
    build_env(
        config, "config.json", args.save_path
    )  # 将a.config复制到a.checkpoint_path/config.json

    logger.info("Initializing Unsupervised Training Process..")

    wav_path2txt(config.data_setting_train.wav_clean_root)
    wav_path2txt(os.path.join(config.data_setting_train.rir_root, "target"))
    wav_path2txt(os.path.join(config.data_setting_train.rir_root, "noisy"))
    wav_path2txt(config.data_setting_val.wav_clean_root)
    wav_path2txt(os.path.join(config.data_setting_val.rir_root, "target"))
    wav_path2txt(os.path.join(config.data_setting_val.rir_root, "noisy"))

    config.data_setting_train.spch_index_txt = os.path.join(
        config.data_setting_train.wav_clean_root, "index.txt"
    )
    config.data_setting_train.rir_index_txt = os.path.join(
        config.data_setting_train.rir_root, "noisy", "index.txt"
    )
    config.data_setting_train.rir_target_index_txt = os.path.join(
        config.data_setting_train.rir_root, "target", "index.txt"
    )
    config.data_setting_val.spch_index_txt = os.path.join(
        config.data_setting_val.wav_clean_root, "index.txt"
    )
    config.data_setting_val.rir_index_txt = os.path.join(
        config.data_setting_val.rir_root, "noisy", "index.txt"
    )
    config.data_setting_val.rir_target_index_txt = os.path.join(
        config.data_setting_val.rir_root, "target", "index.txt"
    )

    if torch.cuda.is_available():
        config.num_gpus = torch.cuda.device_count()
        config.batch_size[0] = int(config.batch_size[0] / config.num_gpus)
        logger.info("Batch size per GPU :" + str(config.batch_size))
    else:
        raise ValueError("No GPU Devices!")

    if config.num_gpus > 1:
        mp.spawn(
            train,
            nprocs=config.num_gpus,
            args=(
                vars(args),
                vars(config),
            ),
        )  # 多线程进行训练，每个GPU一个进程
    else:
        train(0, vars(args), vars(config))


if __name__ == "__main__":
    main()
        
        
        
train_s.py
Audio-WestlakeU/RVAE-EM
import os
import sys
import argparse
import json
import warnings
import random
import numpy as np
import torch
import torch.multiprocessing as mp
import logging
from torch.utils.tensorboard import SummaryWriter
from torch.utils.data import DistributedSampler, DataLoader, RandomSampler
from torch.distributed import init_process_group
from torch.nn.parallel import DistributedDataParallel
from torchinfo import summary
import time
from datetime import datetime
from pesq import pesq
from pystoi import stoi

from utils.env import build_env, dict_to_namespace
from utils.utils import (
    scan_checkpoint,
    load_checkpoint,
    set_optimizer,
    set_scheduler,
    save_checkpoint,
    plot_spectrogram,
    wav_path2txt,
)
from model.RVAE import RVAE as MODEL
from model.lossF import LossF
from dataset.mydataset import MyDataset as dataset
from dataset.io import aud2sptm_woDC as data_preprocess
from dataset.io import sptm_woDC2aud as data_postprocess

# global settings
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
warnings.simplefilter(action="ignore", category=UserWarning)

# CUDNN
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True

logger = logging.getLogger("mylogger")
logger.setLevel(logging.DEBUG)

# training
def train(rank, a, c):

    # create namespace of args and configs
    a = dict_to_namespace(a)
    c = dict_to_namespace(c)

    # logger
    if rank == 0:
        fh = logging.FileHandler(os.path.join(a.save_path, "log.txt"), mode="a")
        fh.setLevel(logging.DEBUG)
        ch = logging.StreamHandler()
        ch.setLevel(logging.DEBUG)
        formatter = logging.Formatter("%(asctime)s %(levelname)s: %(message)s")
        fh.setFormatter(formatter)
        ch.setFormatter(formatter)
        logger.addHandler(fh)
        logger.addHandler(ch)

    # seed
    seed = c.seed
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    seed_this_gpu = c.seed + rank
    torch.cuda.manual_seed(seed_this_gpu)

    # create paths
    os.makedirs(a.save_path, exist_ok=True)
    a.ckpt_save_path = os.path.join(a.save_path, "ckpt")
    os.makedirs(a.ckpt_save_path, exist_ok=True)

    # DDP initialization
    init_process_group(
        backend=c.dist_config.dist_backend,
        init_method=c.dist_config.dist_url,
        world_size=c.dist_config.world_size * c.num_gpus,
        rank=rank,
    )

    # model and optimizer initialization
    device = torch.device("cuda:{:d}".format(rank))
    torch.cuda.set_device(rank)
    torch.cuda.empty_cache()
    model = MODEL(**vars(c.model)).to(device)  # to device
    optim = set_optimizer(model.parameters(), c.optimizer)
    scheduler = set_scheduler(optim, c.lr_scheduler)
    state_dict = None
    steps = 0
    last_epoch = -1
    ckpt = scan_checkpoint(a.ckpt_save_path, "ckpt_")
    if a.start_ckpt:
        if rank == 0:
            logger.info("Loading pretrained model from {}".format(a.start_ckpt))
        state_dict = load_checkpoint(a.start_ckpt, device)
        model.load_state_dict(state_dict["params"])
    else:
        if ckpt:
            if rank == 0:
                logger.info(
                    "Resume training process. Loading checkpoint from {}".format(ckpt)
                )
            state_dict = load_checkpoint(ckpt, device)
            model.load_state_dict(state_dict["params"])
            steps = state_dict["steps"] + 1
            last_epoch = state_dict["epoch"]
            # optim.load_state_dict(state_dict["optim"])
            # scheduler.load_state_dict(state_dict["scheduler"])
        else:
            if rank == 0:
                logger.info("New training process. Saved to {}".format(a.save_path))

    if rank == 0:
        logger.info(summary(model))
    if c.num_gpus > 1:  # DDP model
        model = DistributedDataParallel(
            model,
            device_ids=[rank],
            find_unused_parameters=False,
            broadcast_buffers=False,
        )

    # dataloader
    trainset = dataset(**vars(c.data_setting_train), **vars(c.stft_setting))
    train_sampler = (
        DistributedSampler(trainset) if c.num_gpus > 1 else RandomSampler(trainset)
    )
    train_loader = DataLoader(
        trainset,
        num_workers=c.num_workers,
        shuffle=False,
        sampler=train_sampler,
        batch_size=c.batch_size[0],
        pin_memory=True,
        drop_last=True,
        prefetch_factor=4,
        persistent_workers=True,
    )
    if rank == 0:
        validset = dataset(**vars(c.data_setting_val), **vars(c.stft_setting))
        validation_loader = DataLoader(
            validset,
            num_workers=c.num_workers,
            shuffle=False,
            sampler=None,
            batch_size=c.batch_size[1],
            pin_memory=True,
            drop_last=True,
            prefetch_factor=4,
            persistent_workers=True,
        )
        sw = SummaryWriter(os.path.join(a.save_path, "logs"))

    # training process
    model.train()
    if c.use_amp:
        scaler = torch.cuda.amp.GradScaler()
    for epoch in range(max(0, last_epoch), a.training_epochs):
        if rank == 0:
            start = time.time()
            logger.info("Epoch: {}".format(epoch + 1))
        if c.num_gpus > 1:
            train_sampler.set_epoch(epoch)
        for i, batch in enumerate(train_loader):
            if rank == 0:
                start_b = time.time()
            optim.zero_grad()
            _, audio_clean = batch
            audio_clean = audio_clean.to(device)
            input, _ = data_preprocess(audio_clean, **vars(c.stft_setting))
            if c.use_amp:
                with torch.cuda.amp.autocast():
                    output, z_mean, z_logvar, _ = model(input)
                    loss_IS = LossF().loss_ISD(output, input)
                    loss_KL = LossF().loss_KLD(z_mean, z_logvar)
                    scale_KL = LossF().cal_KL_scale(
                        steps,
                        c.beta,
                        c.beta_zero_step,
                        c.beta_warmup_step,
                        c.beta_holdon_step,
                    )
                    loss_tot = loss_IS + loss_KL * scale_KL
                    assert torch.isnan(loss_tot) == False
                    scaler.scale(loss_tot).backward()
                    if c.gradient_clip:
                        torch.nn.utils.clip_grad.clip_grad_norm(
                            model.parameters(), max_norm=c.gradient_clip
                        )
                    scaler.step(optim)
                    scaler.update()
            else:
                output, z_mean, z_logvar, _ = model(input)
                loss_IS = LossF().loss_ISD(output, input)
                loss_KL = LossF().loss_KLD(z_mean, z_logvar)
                scale_KL = LossF().cal_KL_scale(
                    steps,
                    c.beta,
                    c.beta_zero_step,
                    c.beta_warmup_step,
                    c.beta_holdon_step,
                )
                loss_tot = loss_IS + loss_KL * scale_KL
                assert torch.isnan(loss_tot) == False
                loss_tot.backward()
                if c.gradient_clip:
                    torch.nn.utils.clip_grad.clip_grad_norm(
                        model.parameters(), max_norm=c.gradient_clip
                    )
                optim.step()
            if rank == 0:
                # stdout logger
                if steps % a.stdout_interval == 0:
                    logger.info(
                        "Steps:{:d}, Loss:{:.1f}, Loss_IS:{:.1f}, Loss_KL:{:.1f}, KL_scale:{:.1f}, lr:{:.2e}, time_cost:{:.1f}".format(
                            steps,
                            loss_tot,
                            loss_IS,
                            loss_KL,
                            scale_KL,
                            optim.param_groups[0]["lr"],
                            time.time() - start_b,
                        )
                    )

                # checkpointing
                if steps % a.checkpoint_interval == 0 and steps != 0:
                    checkpoint_filename = "ckpt_{:010d}".format(steps)
                    checkpoint_path = os.path.join(
                        a.ckpt_save_path, checkpoint_filename
                    )
                    save_checkpoint(
                        checkpoint_path,
                        {
                            "params": (
                                model.module if c.num_gpus > 1 else model
                            ).state_dict(),
                            "optim": optim.state_dict(),
                            "scheduler": scheduler.state_dict() if scheduler else None,
                            "steps": steps,
                            "epoch": epoch,
                        },
                    )

                # tensorboard summary logging
                if steps % a.summary_interval == 0:
                    sw.add_scalar("training/loss", loss_tot, steps)
                    sw.add_scalar("training/loss_IS", loss_IS, steps)
                    sw.add_scalar("training/loss_KL", loss_KL, steps)
                    sw.add_scalar("epoch", epoch, steps)
                    sw.add_scalar("lr", optim.param_groups[0]["lr"], steps)

                # validation
                if steps % a.validation_interval == 0:  # and steps != 0:
                    logger.info("validating ...")
                    model.eval()
                    torch.cuda.empty_cache()
                    val_loss_tot_tot = 0
                    val_loss_IS_tot = 0
                    val_loss_KL_tot = 0
                    metric_wbpesq_tot = 0
                    metric_estoi_tot = 0
                    tot_bs = 0
                    with torch.no_grad():
                        for j, batch in enumerate(validation_loader):
                            _, audio_clean = batch
                            audio_clean = audio_clean.to(device)
                            input, input_phase = data_preprocess(
                                audio_clean, **vars(c.stft_setting)
                            )
                            bs = input.shape[0]

                            output, z_mean, z_logvar, _ = model(input)
                            loss_IS = LossF().loss_ISD(output, input)
                            loss_KL = LossF().loss_KLD(z_mean, z_logvar)
                            loss_tot = loss_IS + loss_KL

                            output_wavs = data_postprocess(
                                output, input_phase, **vars(c.stft_setting)
                            )

                            val_loss_tot_tot += loss_tot.item() * bs
                            val_loss_IS_tot += loss_IS.item() * bs
                            val_loss_KL_tot += loss_KL.item() * bs

                            tot_bs += bs

                            if j == 0 and rank == 0:
                                metric_wbpesq_tot = metric_wbpesq_tot + pesq(
                                    c.stft_setting.fs,
                                    audio_clean[0].cpu().numpy(),
                                    output_wavs[0].cpu().numpy(),
                                    "wb",
                                )
                                metric_estoi_tot = metric_estoi_tot + stoi(
                                    audio_clean[0].cpu().numpy(),
                                    output_wavs[0].cpu().numpy(),
                                    c.stft_setting.fs,
                                    extended=True,
                                )
                                spectrogram = (
                                    output[0].permute(1, 0).log().cpu().numpy()
                                )
                                sw.add_figure(
                                    "generated/y_hat_spec_{}".format(steps),
                                    plot_spectrogram(spectrogram),
                                    steps,
                                )
                                if steps == 0:
                                    spectrogram = (
                                        input[0].permute(1, 0).log().cpu().numpy()
                                    )
                                    sw.add_figure(
                                        "clean/y_spec_{}".format(steps),
                                        plot_spectrogram(spectrogram),
                                        steps,
                                    )

                        val_loss_tot = val_loss_tot_tot / tot_bs
                        val_loss_IS = val_loss_IS_tot / tot_bs
                        val_loss_KL = val_loss_KL_tot / tot_bs
                        metric_wbpesq = metric_wbpesq_tot
                        metric_estoi = metric_estoi_tot

                        sw.add_scalar("validation/loss", val_loss_tot, steps)
                        sw.add_scalar("validation/loss_IS", val_loss_IS, steps)
                        sw.add_scalar("validation/loss_KL", val_loss_KL, steps)
                        sw.add_scalar(
                            "validation/metric_wbpesq_1sample", metric_wbpesq, steps
                        )
                        sw.add_scalar(
                            "validation/metric_estoi_1sample", metric_estoi, steps
                        )

                    model.train()

            steps += 1

        if scheduler:
            scheduler.step()

        if rank == 0:
            print(
                "Time taken for epoch {} is {} sec\n".format(
                    epoch + 1, int(time.time() - start)
                )
            )


def main():

    now = datetime.now()
    current_time = now.strftime("%Y-%m-%d-%H-%M")

    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--config", "-c", action="append", default=[], help="config files (.json)"
    )
    parser.add_argument(
        "--save_path", "-p", default="saved_model" + current_time, help="saved path"
    )
    parser.add_argument(
        "--training_epochs", default=100000, type=int, help="maximum epochs"
    )
    parser.add_argument(
        "--stdout_interval", default=10, type=int, help="stdout interval (steps)"
    )
    parser.add_argument(
        "--checkpoint_interval", default=1000, type=int, help="save interval (steps)"
    )
    parser.add_argument(
        "--summary_interval",
        default=100,
        type=int,
        help="tensorboard summary interval (steps)",
    )
    parser.add_argument(
        "--validation_interval",
        default=1000,
        type=int,
        help="validation interval (steps)",
    )
    parser.add_argument(
        "--start_ckpt", "-r", type=str, required=False, help="pretrained model (.ckpt)"
    )

    args = parser.parse_args()

    json_config = {}
    for filename in args.config:
        with open(filename, "r") as f:  # 读取config
            json_config.update(json.load(f))

    config = dict_to_namespace(json_config)  # 转换为属性字典
    build_env(
        config, "config.json", args.save_path
    )  # 将a.config复制到a.checkpoint_path/config.json

    logger.info("Initializing Unsupervised Training Process..")

    wav_path2txt(config.data_setting_train.wav_clean_root)
    wav_path2txt(os.path.join(config.data_setting_train.rir_root, "target"))
    wav_path2txt(os.path.join(config.data_setting_train.rir_root, "noisy"))
    wav_path2txt(config.data_setting_val.wav_clean_root)
    wav_path2txt(os.path.join(config.data_setting_val.rir_root, "target"))
    wav_path2txt(os.path.join(config.data_setting_val.rir_root, "noisy"))

    config.data_setting_train.spch_index_txt = os.path.join(
        config.data_setting_train.wav_clean_root, "index.txt"
    )
    config.data_setting_train.rir_index_txt = os.path.join(
        config.data_setting_train.rir_root, "noisy", "index.txt"
    )
    config.data_setting_train.rir_target_index_txt = os.path.join(
        config.data_setting_train.rir_root, "target", "index.txt"
    )
    config.data_setting_val.spch_index_txt = os.path.join(
        config.data_setting_val.wav_clean_root, "index.txt"
    )
    config.data_setting_val.rir_index_txt = os.path.join(
        config.data_setting_val.rir_root, "noisy", "index.txt"
    )
    config.data_setting_val.rir_target_index_txt = os.path.join(
        config.data_setting_val.rir_root, "target", "index.txt"
    )

    if torch.cuda.is_available():
        config.num_gpus = torch.cuda.device_count()
        config.batch_size[0] = int(config.batch_size[0] / config.num_gpus)
        logger.info("Batch size per GPU :" + str(config.batch_size))
    else:
        raise ValueError("No GPU Devices!")

    if config.num_gpus > 1:
        mp.spawn(
            train,
            nprocs=config.num_gpus,
            args=(
                vars(args),
                vars(config),
            ),
        )  
    else:
        train(0, vars(args), vars(config))


if __name__ == "__main__":
    main()

train_u.py
Audio-WestlakeU/RVAE-EM
import os
import json
import types


class AttrDict(dict):
    def __init__(self, *args, **kwargs):
        super(AttrDict, self).__init__(*args, **kwargs)
        self.__dict__ = self


def dict_to_namespace(d):
    """
    Return: namespace
    Params: dict
    """
    # 如果值是字典，则递归地将其转换为属性字典
    for key, value in d.items():
        if isinstance(value, dict):
            d[key] = dict_to_namespace(value)

    # 将字典转换为属性字典并返回
    return types.SimpleNamespace(**d)


def build_env(config, config_name, path):  # 将配置文件复制到对应目录

    """
    Copy config.json to the folder
    Params:
    config 要保存的字典
    config_name 要保存的文件名
    path 要保存的文件夹
    """
    t_path = os.path.join(path, config_name)
    print("project directory : ", path)
    os.makedirs(path, exist_ok=True)

    json_str = json.dumps(config, cls=NamespaceEncoder)

    with open(t_path, "w") as f:
        f.write(json_str)
    f.close()


class NamespaceEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, types.SimpleNamespace):
            return vars(obj)
        return super().default(obj)

utils/env.py
Audio-WestlakeU/RVAE-EM
import os

import numpy as np

import scipy.stats
from scipy.signal import butter, sosfilt

import torch

from pesq import pesq
from pystoi import stoi


def si_sdr_components(s_hat, s, n):
    """
    """
    # s_target
    alpha_s = np.dot(s_hat, s) / np.linalg.norm(s)**2
    s_target = alpha_s * s

    # e_noise
    alpha_n = np.dot(s_hat, n) / np.linalg.norm(n)**2
    e_noise = alpha_n * n

    # e_art
    e_art = s_hat - s_target - e_noise
    
    return s_target, e_noise, e_art

def energy_ratios(s_hat, s, n):
    """
    """
    s_target, e_noise, e_art = si_sdr_components(s_hat, s, n)

    si_sdr = 10*np.log10(np.linalg.norm(s_target)**2 / np.linalg.norm(e_noise + e_art)**2)
    si_sir = 10*np.log10(np.linalg.norm(s_target)**2 / np.linalg.norm(e_noise)**2)
    si_sar = 10*np.log10(np.linalg.norm(s_target)**2 / np.linalg.norm(e_art)**2)

    return si_sdr, si_sir, si_sar

def mean_conf_int(data, confidence=0.95):
    a = 1.0 * np.array(data)
    n = len(a)
    m, se = np.mean(a), scipy.stats.sem(a)
    h = se * scipy.stats.t.ppf((1 + confidence) / 2., n-1)
    return m, h

class Method():
    def __init__(self, name, base_dir, metrics):
        self.name = name
        self.base_dir = base_dir
        self.metrics = {} 
        
        for i in range(len(metrics)):
            metric = metrics[i]
            value = []
            self.metrics[metric] = value 
            
    def append(self, matric, value):
        self.metrics[matric].append(value)

    def get_mean_ci(self, metric):
        return mean_conf_int(np.array(self.metrics[metric]))

def hp_filter(signal, cut_off=80, order=10, sr=16000):
    factor = cut_off /sr * 2
    sos = butter(order, factor, 'hp', output='sos')
    filtered = sosfilt(sos, signal)
    return filtered

def si_sdr(s, s_hat):
    alpha = np.dot(s_hat, s)/np.linalg.norm(s)**2   
    sdr = 10*np.log10(np.linalg.norm(alpha*s)**2/np.linalg.norm(
        alpha*s - s_hat)**2)
    return sdr

def snr_dB(s,n):
    s_power = 1/len(s)*np.sum(s**2)
    n_power = 1/len(n)*np.sum(n**2)
    snr_dB = 10*np.log10(s_power/n_power)
    return snr_dB

def pad_spec(Y):
    T = Y.size(3)
    if T%64 !=0:
        num_pad = 64-T%64
    else:
        num_pad = 0
    pad2d = torch.nn.ZeroPad2d((0, num_pad, 0,0))
    return pad2d(Y)


def ensure_dir(file_path):
    directory = file_path
    if not os.path.exists(directory):
        os.makedirs(directory)


def print_metrics(x, y, x_hat_list, labels, sr=16000):
    _si_sdr_mix = si_sdr(x, y)
    _pesq_mix = pesq(sr, x, y, 'wb')
    _estoi_mix = stoi(x, y, sr, extended=True)
    print(f'Mixture:  PESQ: {_pesq_mix:.2f}, ESTOI: {_estoi_mix:.2f}, SI-SDR: {_si_sdr_mix:.2f}')
    for i, x_hat in enumerate(x_hat_list):
        _si_sdr = si_sdr(x, x_hat)
        _pesq = pesq(sr, x, x_hat, 'wb')
        _estoi = stoi(x, x_hat, sr, extended=True)
        print(f'{labels[i]}: {_pesq:.2f}, ESTOI: {_estoi:.2f}, SI-SDR: {_si_sdr:.2f}')

def mean_std(data):
    # data = data[~np.isnan(data)]
    mean = np.mean(data)
    std = np.std(data)
    return mean, std

def print_mean_std(data, decimal=2):
    data = np.array(data)
    # data = data[~np.isnan(data)]
    mean = np.mean(data)
    std = np.std(data)
    if decimal == 2:
        string = f'{mean:.2f} ± {std:.2f}'
    elif decimal == 1:
        string = f'{mean:.1f} ± {std:.1f}'
    return string

utils/other.py
Audio-WestlakeU/RVAE-EM
import os
from glob import glob
import torch
import librosa
import numpy as np
import soundfile as sf
import matplotlib.pyplot as plt
import yaml
import sys


def scan_checkpoint(cp_dir, prefix=""):
    pattern = os.path.join(cp_dir, prefix + "??????????")
    cp_list = glob(pattern)
    if len(cp_list) == 0:
        return None
    return sorted(cp_list)[-1]


def load_checkpoint(filepath, device):
    assert os.path.isfile(filepath)
    checkpoint_dict = torch.load(filepath, map_location="cpu")
    return checkpoint_dict


def set_optimizer(params, config_optim):

    """
    设置优化器
    :params
    params 优化的目标
    config_optim: 配置
    :ret
    optimizer
    """

    if config_optim.type.upper() == "SGD":
        optim = torch.optim.SGD(params, **vars(config_optim.kwargs))
    elif config_optim.type.upper() == "ADAM":
        optim = torch.optim.Adam(params, **vars(config_optim.kwargs))
    elif config_optim.type.upper() == "ADAMW":
        optim = torch.optim.AdamW(params, **vars(config_optim.kwargs))
    else:
        raise ValueError("Wrong Optimizer Type")
    return optim


def set_scheduler(optim, config_lrscheduler):
    if config_lrscheduler.type.upper() == "ReduceLROnPlateau".upper():
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optim, **vars(config_lrscheduler.kwargs)
        )
    elif config_lrscheduler.type.upper() == "none".upper():
        scheduler = None
    elif config_lrscheduler.type.upper() == "ExponentialLR".upper():
        scheduler = torch.optim.lr_scheduler.ExponentialLR(
            optim, **vars(config_lrscheduler.kwargs)
        )
    elif config_lrscheduler.type.upper() == "CyclicLR".upper():
        scheduler = torch.optim.lr_scheduler.CyclicLR(
            optim, **vars(config_lrscheduler.kwargs), cycle_momentum=False
        )
    elif config_lrscheduler.type.upper() == "StepLR".upper():
        scheduler = torch.optim.lr_scheduler.StepLR(
            optim, **vars(config_lrscheduler.kwargs)
        )
    else:
        raise ValueError("Wrong Scheduler Type")
    return scheduler


def get_filelist_from_txt(filepath):
    with open(filepath, "r", encoding="utf-8") as f:
        files = [x for x in f.read().split("\n") if len(x) > 0]
    return files


def save_checkpoint(filepath, obj):
    print("Saving checkpoint to {}".format(filepath))
    torch.save(obj, filepath)


def plot_spectrogram(spectrogram):
    spectrogram[spectrogram < -8] = -8
    fig = plt.figure()
    plt.imshow(spectrogram, aspect="auto", origin="lower")
    plt.colorbar()
    plt.close()

    return fig


def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


def add_DC_and_save_wav(
    spectrogram_no_DC, phase_no_DC, stft_config, dir_path, filename
):
    spectrogram_no_DC = spectrogram_no_DC.detach()
    phase_no_DC = phase_no_DC.detach()
    device = spectrogram_no_DC.device
    bs, F, T = spectrogram_no_DC.shape
    F = F + 1
    spectrogram = torch.zeros([bs, F, T]).to(device) + 0j
    spectrogram[:, 1:, :] = spectrogram_no_DC
    spectrogram = spectrogram.sqrt()
    spectrogram = np.array(spectrogram.to("cpu"))
    phase = torch.zeros([bs, F, T]).to(device)
    phase[:, 1:, :] = phase_no_DC
    phase = torch.exp(1j * phase)
    phase = np.array(phase.to("cpu"))
    spectrogram = spectrogram * phase

    for isample in range(bs):
        wav = librosa.istft(
            spectrogram[isample],
            hop_length=stft_config.hop,
            win_length=stft_config.nfft,
            n_fft=stft_config.nfft,
            window=stft_config.win,
        )
        wav = wav / np.max(np.abs(wav))
        # wav = np.expand_dims(wav,0)
        filepath = os.path.join(dir_path, filename[isample])
        sf.write(filepath, wav, 16000)

    return None


def del_files_from_dir(filepath):
    del_list = os.listdir(filepath)
    for f in del_list:
        file_path = os.path.join(filepath, f)
        if os.path.isfile(file_path):
            os.remove(file_path)
    return None


def read_yaml(config_path):
    # 读取yaml文件
    with open(config_path, encoding="utf-8") as f:
        config = yaml.load(f.read(), Loader=yaml.FullLoader)
    return config


def wav_path2txt(root):
    """
    Function: save the path of .wavs in the root folder to index.txt
    Params:
        root: root folder
    """
    # 获取文件夹中wav文件的地址存入txt文件
    files = sorted(glob("{}/**/*.wav".format(root), recursive=True))
    f = open(os.path.join(root, "index.txt"), "w")
    for filename in files:
        f.write(filename + os.linesep)
    f.close()
    return None


class TxtLogger(object):
    def __init__(self, filename="Default.log"):
        self.terminal = sys.stdout
        self.log = open(filename, "a")

    def write(self, message):
        self.terminal.write(message)
        self.log.write(message)

    def flush(self):
        pass

utils/utils.py
Audio-WestlakeU/RVAE-EM