Spaces:
Runtime error
Runtime error
| import math | |
| from typing import Optional | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import torchaudio | |
| from torchaudio.functional import resample | |
| from huggingface_hub import PyTorchModelHubMixin | |
| class SEModule(nn.Module): | |
| def __init__(self, channels : int , bottleneck : int = 128) -> None: | |
| super(SEModule, self).__init__() | |
| self.se = nn.Sequential( | |
| nn.AdaptiveAvgPool1d(1), | |
| nn.Conv1d(channels, bottleneck, kernel_size=1, padding=0), | |
| nn.ReLU(), | |
| # nn.BatchNorm1d(bottleneck), # I remove this layer | |
| nn.Conv1d(bottleneck, channels, kernel_size=1, padding=0), | |
| nn.Sigmoid(), | |
| ) | |
| def forward(self, input : torch.Tensor) -> torch.Tensor: | |
| x = self.se(input) | |
| return input * x | |
| class Bottle2neck(nn.Module): | |
| def __init__(self, inplanes : int, planes : int, kernel_size : Optional[int] = None, dilation : Optional[int] = None, scale : int = 8) -> None: | |
| super(Bottle2neck, self).__init__() | |
| width = int(math.floor(planes / scale)) | |
| self.conv1 = nn.Conv1d(inplanes, width*scale, kernel_size=1) | |
| self.bn1 = nn.BatchNorm1d(width*scale) | |
| self.nums = scale -1 | |
| convs = [] | |
| bns = [] | |
| num_pad = math.floor(kernel_size/2)*dilation | |
| for i in range(self.nums): | |
| convs.append(nn.Conv1d(width, width, kernel_size=kernel_size, dilation=dilation, padding=num_pad)) | |
| bns.append(nn.BatchNorm1d(width)) | |
| self.convs = nn.ModuleList(convs) | |
| self.bns = nn.ModuleList(bns) | |
| self.conv3 = nn.Conv1d(width*scale, planes, kernel_size=1) | |
| self.bn3 = nn.BatchNorm1d(planes) | |
| self.relu = nn.ReLU() | |
| self.width = width | |
| self.se = SEModule(planes) | |
| def forward(self, x : torch.Tensor) -> torch.Tensor: | |
| residual = x | |
| out = self.conv1(x) | |
| out = self.relu(out) | |
| out = self.bn1(out) | |
| spx = torch.split(out, self.width, 1) | |
| for i in range(self.nums): | |
| if i==0: | |
| sp = spx[i] | |
| else: | |
| sp = sp + spx[i] | |
| sp = self.convs[i](sp) | |
| sp = self.relu(sp) | |
| sp = self.bns[i](sp) | |
| if i==0: | |
| out = sp | |
| else: | |
| out = torch.cat((out, sp), 1) | |
| out = torch.cat((out, spx[self.nums]),1) | |
| out = self.conv3(out) | |
| out = self.relu(out) | |
| out = self.bn3(out) | |
| out = self.se(out) | |
| out += residual | |
| return out | |
| class ECAPA_gender(nn.Module, PyTorchModelHubMixin): | |
| def __init__(self, C : int = 1024): | |
| super(ECAPA_gender, self).__init__() | |
| self.C = C | |
| self.conv1 = nn.Conv1d(80, C, kernel_size=5, stride=1, padding=2) | |
| self.relu = nn.ReLU() | |
| self.bn1 = nn.BatchNorm1d(C) | |
| self.layer1 = Bottle2neck(C, C, kernel_size=3, dilation=2, scale=8) | |
| self.layer2 = Bottle2neck(C, C, kernel_size=3, dilation=3, scale=8) | |
| self.layer3 = Bottle2neck(C, C, kernel_size=3, dilation=4, scale=8) | |
| # I fixed the shape of the output from MFA layer, that is close to the setting from ECAPA paper. | |
| self.layer4 = nn.Conv1d(3*C, 1536, kernel_size=1) | |
| self.attention = nn.Sequential( | |
| nn.Conv1d(4608, 256, kernel_size=1), | |
| nn.ReLU(), | |
| nn.BatchNorm1d(256), | |
| nn.Tanh(), # I add this layer | |
| nn.Conv1d(256, 1536, kernel_size=1), | |
| nn.Softmax(dim=2), | |
| ) | |
| self.bn5 = nn.BatchNorm1d(3072) | |
| self.fc6 = nn.Linear(3072, 192) | |
| self.bn6 = nn.BatchNorm1d(192) | |
| self.fc7 = nn.Linear(192, 2) | |
| self.pred2gender = {0 : 'male', 1 : 'female'} | |
| def logtorchfbank(self, x : torch.Tensor) -> torch.Tensor: | |
| # Preemphasis | |
| flipped_filter = torch.FloatTensor([-0.97, 1.]).unsqueeze(0).unsqueeze(0) | |
| x = x.unsqueeze(1) | |
| x = F.pad(x, (1, 0), 'reflect') | |
| x = F.conv1d(x, flipped_filter).squeeze(1) | |
| # Melspectrogram | |
| x = torchaudio.transforms.MelSpectrogram(sample_rate=16000, n_fft=512, win_length=400, hop_length=160, \ | |
| f_min = 20, f_max = 7600, window_fn=torch.hamming_window, n_mels=80)(x) + 1e-6 | |
| # Log and normalize | |
| x = x.log() | |
| x = x - torch.mean(x, dim=-1, keepdim=True) | |
| return x | |
| def forward(self, x : torch.Tensor) -> torch.Tensor: | |
| x = self.logtorchfbank(x) | |
| x = self.conv1(x) | |
| x = self.relu(x) | |
| x = self.bn1(x) | |
| x1 = self.layer1(x) | |
| x2 = self.layer2(x+x1) | |
| x3 = self.layer3(x+x1+x2) | |
| x = self.layer4(torch.cat((x1,x2,x3),dim=1)) | |
| x = self.relu(x) | |
| t = x.size()[-1] | |
| global_x = torch.cat((x,torch.mean(x,dim=2,keepdim=True).repeat(1,1,t), torch.sqrt(torch.var(x,dim=2,keepdim=True).clamp(min=1e-4)).repeat(1,1,t)), dim=1) | |
| w = self.attention(global_x) | |
| mu = torch.sum(x * w, dim=2) | |
| sg = torch.sqrt( ( torch.sum((x**2) * w, dim=2) - mu**2 ).clamp(min=1e-4) ) | |
| x = torch.cat((mu,sg),1) | |
| x = self.bn5(x) | |
| x = self.fc6(x) | |
| x = self.bn6(x) | |
| x = self.relu(x) | |
| x = self.fc7(x) | |
| return x | |
| def load_audio(self, path : str) -> torch.Tensor: | |
| audio, sr = torchaudio.load(path) | |
| if sr != 16000: | |
| audio = resample(audio, sr, 16000) | |
| return audio | |
| def predict(self, audio : torch.Tensor) -> torch.Tensor: | |
| audio = self.load_audio(audio) | |
| self.eval() | |
| with torch.no_grad(): | |
| output = self.forward(audio) | |
| _, pred = output.max(1) | |
| return self.pred2gender[pred.item()] |