(self, x1, x2,lens=None)
| 143 | self.enh_block = DPCRN(feat_dim=self.freq_dim,input_channel_rate=2) |
| 144 | |
| 145 | def forward(self, x1, x2,lens=None): |
| 146 | |
| 147 | specs1 = self.stft(x1) |
| 148 | real1 = specs1[:,:self.fft_len//2+1] |
| 149 | imag1 = specs1[:,self.fft_len//2+1:] |
| 150 | spec_mags1 = torch.sqrt(real1**2+imag1**2+1e-8) |
| 151 | spec_phase1 = torch.atan2(imag1, real1) |
| 152 | |
| 153 | specs2 = self.stft(x2) |
| 154 | real2 = specs2[:,:self.fft_len//2+1] |
| 155 | imag2 = specs2[:,self.fft_len//2+1:] |
| 156 | spec_mags2 = torch.sqrt(real2**2+imag2**2+1e-8) |
| 157 | spec_phase2 = torch.atan2(imag2, real2) |
| 158 | |
| 159 | real = torch.stack([real1,real2],1) |
| 160 | imag = torch.stack([imag1,imag2],1) |
| 161 | cspecs = torch.cat([real,imag],1) |
| 162 | cspecs = cspecs[:,:,1:self.freq_dim+1] |
| 163 | |
| 164 | |
| 165 | out = cspecs#self.cln(cspecs) |
| 166 | out = self.enh_block(out) |
| 167 | |
| 168 | mask_real = out[:,0] |
| 169 | mask_imag = out[:,1] |
| 170 | mask_real = F.pad(mask_real, [0,0,1,0]) |
| 171 | mask_imag = F.pad(mask_imag, [0,0,1,0]) |
| 172 | |
| 173 | spec_mags2 = spec_mags2[:,:self.freq_dim+1,:] |
| 174 | spec_phase2 = spec_phase2[:,:self.freq_dim+1,:] |
| 175 | real2 = real2[:,:self.freq_dim+1,:] |
| 176 | imag2 = imag2[:,:self.freq_dim+1,:] |
| 177 | |
| 178 | if self.masking_mode == 'E' : |
| 179 | mask_mags = (mask_real**2+mask_imag**2)**0.5 |
| 180 | real_phase = mask_real/(mask_mags+1e-8) |
| 181 | imag_phase = mask_imag/(mask_mags+1e-8) |
| 182 | mask_phase = torch.atan2( |
| 183 | imag_phase, |
| 184 | real_phase |
| 185 | ) |
| 186 | |
| 187 | #mask_mags = torch.clamp_(mask_mags,0,100) |
| 188 | mask_mags = torch.tanh(mask_mags) |
| 189 | est_mags = mask_mags*spec_mags2 |
| 190 | est_phase = spec_phase2 + mask_phase |
| 191 | real = est_mags*torch.cos(est_phase) |
| 192 | imag = est_mags*torch.sin(est_phase) |
| 193 | elif self.masking_mode == 'C': |
| 194 | real,imag = real2*mask_real-imag2*mask_imag + real1[:,:self.freq_dim+1,:], real2*mask_imag+imag2*mask_real + imag1[:,:self.freq_dim+1,:] |
| 195 | |
| 196 | real = torch.cat([real, real1[:,self.freq_dim+1:,:]],1) |
| 197 | imag = torch.cat([imag, imag1[:,self.freq_dim+1:,:]],1) |
| 198 | out_spec = torch.cat([real, imag], 1) |
| 199 | out_wav = self.istft(out_spec) |
| 200 | |
| 201 | out_wav = torch.squeeze(out_wav, 1) |
| 202 | #out_wav = torch.tanh(out_wav) |
nothing calls this directly
no outgoing calls
no test coverage detected