Spaces:
Running
Running
File size: 3,255 Bytes
f74ae8e a813eb0 f74ae8e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
#!/usr/bin/python3
# -*- coding: utf-8 -*-
from einops.layers.torch import Rearrange
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from pesq import pesq
from joblib import Parallel, delayed
def phase_losses(phase_r, phase_g):
ip_loss = torch.mean(anti_wrapping_function(phase_r - phase_g))
gd_loss = torch.mean(anti_wrapping_function(torch.diff(phase_r, dim=1) - torch.diff(phase_g, dim=1)))
iaf_loss = torch.mean(anti_wrapping_function(torch.diff(phase_r, dim=2) - torch.diff(phase_g, dim=2)))
return ip_loss, gd_loss, iaf_loss
def anti_wrapping_function(x):
return torch.abs(x - torch.round(x / (2 * np.pi)) * 2 * np.pi)
def pesq_score(utts_r, utts_g, h):
pesq_score = Parallel(n_jobs=30)(delayed(eval_pesq)(
utts_r[i].squeeze().cpu().numpy(),
utts_g[i].squeeze().cpu().numpy(),
h.sample_rate)
for i in range(len(utts_r)))
pesq_score = np.mean(pesq_score)
return pesq_score
def eval_pesq(clean_utt, esti_utt, sr):
try:
pesq_score = pesq(sr, clean_utt, esti_utt)
except:
pesq_score = -1
return pesq_score
def mag_pha_stft(y, n_fft, hop_size, win_size, compress_factor=1.0, center=True):
hann_window = torch.hann_window(win_size).to(y.device)
stft_spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window,
center=center, pad_mode='reflect', normalized=False, return_complex=True)
stft_spec = torch.view_as_real(stft_spec)
mag = torch.sqrt(stft_spec.pow(2).sum(-1) + 1e-9)
pha = torch.atan2(stft_spec[:, :, :, 1] + 1e-10, stft_spec[:, :, :, 0] + 1e-5)
# Magnitude Compression
mag = torch.pow(mag, compress_factor)
com = torch.stack((mag*torch.cos(pha), mag*torch.sin(pha)), dim=-1)
return mag, pha, com
def mag_pha_istft(mag, pha, n_fft, hop_size, win_size, compress_factor=1.0, center=True):
# Magnitude Decompression
mag = torch.pow(mag, (1.0/compress_factor))
com = torch.complex(mag*torch.cos(pha), mag*torch.sin(pha))
hann_window = torch.hann_window(win_size).to(com.device)
wav = torch.istft(com, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window, center=center)
return wav
class LearnableSigmoid1d(nn.Module):
def __init__(self, in_features, beta=1):
super().__init__()
self.beta = beta
self.slope = nn.Parameter(torch.ones(in_features))
self.slope.requiresGrad = True
def forward(self, x):
# x shape: [batch_size, time_steps, spec_bins]
return self.beta * torch.sigmoid(self.slope * x)
class LearnableSigmoid2d(nn.Module):
def __init__(self, in_features, beta=1):
super().__init__()
self.beta = beta
self.slope = nn.Parameter(torch.ones(in_features, 1))
self.slope.requiresGrad = True
def forward(self, x):
return self.beta * torch.sigmoid(self.slope * x)
def main():
learnable_sigmoid = LearnableSigmoid1d(201)
a = torch.randn(4, 100, 201)
result = learnable_sigmoid.forward(a)
print(result.shape)
return
if __name__ == '__main__':
main()
|