File size: 1,326 Bytes
28e6f98 | 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 | import torch
def add_frequency_noise(fft, snr=10, vacant_snr=15, mask=None):
### 根据SNR确定noise的放大比例
num_pixels = fft.numel()
fft_magnitude = torch.abs(fft)
fft_phase = torch.angle(fft)
# fft_magnitude
mag_psr = torch.mean(torch.abs(fft_magnitude) ** 2)
mag_pnr = mag_psr / (10 ** (snr / 10)) # Calculate noise power
noise_mag = torch.randn_like(fft_magnitude) * torch.sqrt(mag_pnr)
mag_psr_vacant = mag_psr / (10 ** (vacant_snr / 10))
noise_mag_vacant = torch.randn_like(fft_magnitude) * torch.sqrt(mag_psr_vacant)
fft_magnitude = fft_magnitude + \
noise_mag * fft_magnitude * mask + \
noise_mag_vacant * (1- mask)
fft_magnitude = torch.abs(fft_magnitude)
# fft_phase
pha_psr = torch.mean(torch.abs(fft_phase) ** 2)
pha_pnr = pha_psr / (10 ** (snr / 10)) # Calculate noise power for phase
noise_pha = torch.randn_like(fft_phase) * torch.sqrt(pha_pnr)
pha_psr_vacant = pha_psr / (10 ** (vacant_snr / 10))
noise_pha_vacant = torch.randn_like(fft_phase) * torch.sqrt(pha_psr_vacant)
fft_phase = fft_phase + \
noise_pha * fft_phase * mask + \
noise_pha_vacant * (1- mask)
noise_fft = fft_magnitude * torch.exp(1j * fft_phase)
return noise_fft
|