Create utils/noisereduce.py

tools/utils/noisereduce.py

@@ -0,0 +1,196 @@
+import torch
+import tempfile
+import numpy as np
+
+from joblib import Parallel, delayed
+from torch.nn.functional import conv1d, conv2d
+
+@torch.no_grad()
+def amp_to_db(x, eps = torch.finfo(torch.float32).eps, top_db = 40):
+    x_db = 20 * torch.log10(x.abs() + eps)
+    return torch.max(x_db, (x_db.max(-1).values - top_db).unsqueeze(-1))
+
+@torch.no_grad()
+def temperature_sigmoid(x, x0, temp_coeff):
+    return torch.sigmoid((x - x0) / temp_coeff)
+
+@torch.no_grad()
+def linspace(start, stop, num = 50, endpoint = True, **kwargs):
+    return torch.linspace(start, stop, num, **kwargs) if endpoint else torch.linspace(start, stop, num + 1, **kwargs)[:-1]
+
+def _smoothing_filter(n_grad_freq, n_grad_time):
+    smoothing_filter = np.outer(np.concatenate([np.linspace(0, 1, n_grad_freq + 1, endpoint=False), np.linspace(1, 0, n_grad_freq + 2)])[1:-1], np.concatenate([np.linspace(0, 1, n_grad_time + 1, endpoint=False), np.linspace(1, 0, n_grad_time + 2)])[1:-1])
+    return smoothing_filter / np.sum(smoothing_filter)
+
+class SpectralGate:
+    def __init__(self, y, sr, prop_decrease, chunk_size, padding, n_fft, win_length, hop_length, time_constant_s, freq_mask_smooth_hz, time_mask_smooth_ms, tmp_folder, use_tqdm, n_jobs):
+        self.sr = sr
+        self.flat = False
+        y = np.array(y)
+
+        if len(y.shape) == 1:
+            self.y = np.expand_dims(y, 0)
+            self.flat = True
+        elif len(y.shape) > 2: raise ValueError
+        else: self.y = y
+
+        self._dtype = y.dtype
+        self.n_channels, self.n_frames = self.y.shape
+        self._chunk_size = chunk_size
+        self.padding = padding
+        self.n_jobs = n_jobs
+        self.use_tqdm = use_tqdm
+        self._tmp_folder = tmp_folder
+        self._n_fft = n_fft
+        self._win_length = self._n_fft if win_length is None else win_length
+        self._hop_length = (self._win_length // 4) if hop_length is None else hop_length
+        self._time_constant_s = time_constant_s
+        self._prop_decrease = prop_decrease
+
+        if (freq_mask_smooth_hz is None) & (time_mask_smooth_ms is None): self.smooth_mask = False
+        else: self._generate_mask_smoothing_filter(freq_mask_smooth_hz, time_mask_smooth_ms)
+
+    def _generate_mask_smoothing_filter(self, freq_mask_smooth_hz, time_mask_smooth_ms):
+        if freq_mask_smooth_hz is None: n_grad_freq = 1
+        else:
+            n_grad_freq = int(freq_mask_smooth_hz / (self.sr / (self._n_fft / 2)))
+            if n_grad_freq < 1: raise ValueError
+
+        if time_mask_smooth_ms is None: n_grad_time = 1
+        else:
+            n_grad_time = int(time_mask_smooth_ms / ((self._hop_length / self.sr) * 1000))
+            if n_grad_time < 1: raise ValueError
+
+        if (n_grad_time == 1) & (n_grad_freq == 1): self.smooth_mask = False
+        else:
+            self.smooth_mask = True
+            self._smoothing_filter = _smoothing_filter(n_grad_freq, n_grad_time)
+
+    def _read_chunk(self, i1, i2):
+        i1b = 0 if i1 < 0 else i1  
+        i2b = self.n_frames if i2 > self.n_frames else i2 
+        chunk = np.zeros((self.n_channels, i2 - i1))
+        chunk[:, i1b - i1: i2b - i1] = self.y[:, i1b:i2b]
+        return chunk
+
+    def filter_chunk(self, start_frame, end_frame):
+        i1 = start_frame - self.padding
+        return self._do_filter(self._read_chunk(i1, (end_frame + self.padding)))[:, start_frame - i1: end_frame - i1]
+
+    def _get_filtered_chunk(self, ind):
+        start0 = ind * self._chunk_size
+        end0 = (ind + 1) * self._chunk_size
+        return self.filter_chunk(start_frame=start0, end_frame=end0)
+
+    def _do_filter(self, chunk):
+        pass
+
+    def _iterate_chunk(self, filtered_chunk, pos, end0, start0, ich):
+        filtered_chunk[:, pos: pos + end0 - start0] = self._get_filtered_chunk(ich)[:, start0:end0]
+        pos += end0 - start0
+
+    def get_traces(self, start_frame=None, end_frame=None):
+        if start_frame is None: start_frame = 0
+        if end_frame is None: end_frame = self.n_frames
+
+        if self._chunk_size is not None:
+            if end_frame - start_frame > self._chunk_size:
+                ich1 = int(start_frame / self._chunk_size)
+                ich2 = int((end_frame - 1) / self._chunk_size)
+
+                with tempfile.NamedTemporaryFile(prefix=self._tmp_folder) as fp:
+                    filtered_chunk = np.memmap(fp, dtype=self._dtype, shape=(self.n_channels, int(end_frame - start_frame)), mode="w+")
+                    pos_list, start_list, end_list = [], [], []
+                    pos = 0
+
+                    for ich in range(ich1, ich2 + 1):
+                        start0 = (start_frame - ich * self._chunk_size) if ich == ich1 else 0
+                        end0 = end_frame - ich * self._chunk_size if ich == ich2 else self._chunk_size
+                        pos_list.append(pos)
+                        start_list.append(start0)
+                        end_list.append(end0)
+                        pos += end0 - start0
+
+                    Parallel(n_jobs=self.n_jobs)(delayed(self._iterate_chunk)(filtered_chunk, pos, end0, start0, ich) for pos, start0, end0, ich in zip(pos_list, start_list, end_list, range(ich1, ich2 + 1)))
+                    return filtered_chunk.astype(self._dtype).flatten() if self.flat else filtered_chunk.astype(self._dtype)
+
+        filtered_chunk = self.filter_chunk(start_frame=0, end_frame=end_frame)
+        return filtered_chunk.astype(self._dtype).flatten() if self.flat else filtered_chunk.astype(self._dtype)
+
+class TG(torch.nn.Module):
+    @torch.no_grad()
+    def __init__(self, sr, nonstationary = False, n_std_thresh_stationary = 1.5, n_thresh_nonstationary = 1.3, temp_coeff_nonstationary = 0.1, n_movemean_nonstationary = 20, prop_decrease = 1.0, n_fft = 1024, win_length = None, hop_length = None, freq_mask_smooth_hz = 500, time_mask_smooth_ms = 50):
+        super().__init__()
+        self.sr = sr
+        self.nonstationary = nonstationary
+        assert 0.0 <= prop_decrease <= 1.0
+        self.prop_decrease = prop_decrease
+        self.n_fft = n_fft
+        self.win_length = self.n_fft if win_length is None else win_length
+        self.hop_length = self.win_length // 4 if hop_length is None else hop_length
+        self.n_std_thresh_stationary = n_std_thresh_stationary
+        self.temp_coeff_nonstationary = temp_coeff_nonstationary
+        self.n_movemean_nonstationary = n_movemean_nonstationary
+        self.n_thresh_nonstationary = n_thresh_nonstationary
+        self.freq_mask_smooth_hz = freq_mask_smooth_hz
+        self.time_mask_smooth_ms = time_mask_smooth_ms
+        self.register_buffer("smoothing_filter", self._generate_mask_smoothing_filter())
+
+    @torch.no_grad()
+    def _generate_mask_smoothing_filter(self):
+        if self.freq_mask_smooth_hz is None and self.time_mask_smooth_ms is None: return None
+        n_grad_freq = (1 if self.freq_mask_smooth_hz is None else int(self.freq_mask_smooth_hz / (self.sr / (self.n_fft / 2))))
+        if n_grad_freq < 1: raise ValueError
+
+        n_grad_time = (1 if self.time_mask_smooth_ms is None else int(self.time_mask_smooth_ms / ((self.hop_length / self.sr) * 1000)))
+        if n_grad_time < 1: raise ValueError
+        if n_grad_time == 1 and n_grad_freq == 1: return None
+
+        smoothing_filter = torch.outer(torch.cat([linspace(0, 1, n_grad_freq + 1, endpoint=False), linspace(1, 0, n_grad_freq + 2)])[1:-1], torch.cat([linspace(0, 1, n_grad_time + 1, endpoint=False), linspace(1, 0, n_grad_time + 2)])[1:-1]).unsqueeze(0).unsqueeze(0)
+        return smoothing_filter / smoothing_filter.sum()
+
+    @torch.no_grad()
+    def _stationary_mask(self, X_db, xn = None):
+        XN_db = amp_to_db(torch.stft(xn, n_fft=self.n_fft, hop_length=self.hop_length, win_length=self.win_length, return_complex=True, pad_mode="constant", center=True, window=torch.hann_window(self.win_length).to(xn.device))).to(dtype=X_db.dtype) if xn is not None else X_db
+        std_freq_noise, mean_freq_noise = torch.std_mean(XN_db, dim=-1)
+        return torch.gt(X_db, (mean_freq_noise + std_freq_noise * self.n_std_thresh_stationary).unsqueeze(2))
+
+    @torch.no_grad()
+    def _nonstationary_mask(self, X_abs):
+        X_smoothed = (conv1d(X_abs.reshape(-1, 1, X_abs.shape[-1]), torch.ones(self.n_movemean_nonstationary, dtype=X_abs.dtype, device=X_abs.device).view(1, 1, -1), padding="same").view(X_abs.shape) / self.n_movemean_nonstationary)
+        return temperature_sigmoid(((X_abs - X_smoothed) / X_smoothed), self.n_thresh_nonstationary, self.temp_coeff_nonstationary)
+
+    def forward(self, x, xn = None):
+        assert x.ndim == 2
+        if x.shape[-1] < self.win_length * 2: raise Exception
+        assert xn is None or xn.ndim == 1 or xn.ndim == 2
+        if xn is not None and xn.shape[-1] < self.win_length * 2: raise Exception
+
+        X = torch.stft(x, n_fft=self.n_fft, hop_length=self.hop_length, win_length=self.win_length, return_complex=True, pad_mode="constant", center=True, window=torch.hann_window(self.win_length).to(x.device))
+        sig_mask = self._nonstationary_mask(X.abs()) if self.nonstationary else self._stationary_mask(amp_to_db(X), xn)
+
+        sig_mask = self.prop_decrease * (sig_mask * 1.0 - 1.0) + 1.0
+        if self.smoothing_filter is not None: sig_mask = conv2d(sig_mask.unsqueeze(1), self.smoothing_filter.to(sig_mask.dtype), padding="same")
+
+        Y = X * sig_mask.squeeze(1)
+        return torch.istft(Y, n_fft=self.n_fft, hop_length=self.hop_length, win_length=self.win_length, center=True, window=torch.hann_window(self.win_length).to(Y.device)).to(dtype=x.dtype)
+
+class StreamedTorchGate(SpectralGate):
+    def __init__(self, y, sr, stationary=False, y_noise=None, prop_decrease=1.0, time_constant_s=2.0, freq_mask_smooth_hz=500, time_mask_smooth_ms=50, thresh_n_mult_nonstationary=2, sigmoid_slope_nonstationary=10, n_std_thresh_stationary=1.5, tmp_folder=None, chunk_size=600000, padding=30000, n_fft=1024, win_length=None, hop_length=None, clip_noise_stationary=True, use_tqdm=False, n_jobs=1, device="cpu"):
+        super().__init__(y=y, sr=sr, chunk_size=chunk_size, padding=padding, n_fft=n_fft, win_length=win_length, hop_length=hop_length, time_constant_s=time_constant_s, freq_mask_smooth_hz=freq_mask_smooth_hz, time_mask_smooth_ms=time_mask_smooth_ms, tmp_folder=tmp_folder, prop_decrease=prop_decrease, use_tqdm=use_tqdm, n_jobs=n_jobs)
+        self.device = torch.device(device if torch.cuda.is_available() else 'cpu')
+
+        if y_noise is not None:
+            if y_noise.shape[-1] > y.shape[-1] and clip_noise_stationary: y_noise = y_noise[: y.shape[-1]]
+            y_noise = torch.from_numpy(y_noise).to(device)
+            if len(y_noise.shape) == 1: y_noise = y_noise.unsqueeze(0)
+
+        self.y_noise = y_noise
+        self.tg = TG(sr=sr, nonstationary=not stationary, n_std_thresh_stationary=n_std_thresh_stationary, n_thresh_nonstationary=thresh_n_mult_nonstationary, temp_coeff_nonstationary=1 / sigmoid_slope_nonstationary, n_movemean_nonstationary=int(time_constant_s / self._hop_length * sr), prop_decrease=prop_decrease, n_fft=self._n_fft, win_length=self._win_length, hop_length=self._hop_length, freq_mask_smooth_hz=freq_mask_smooth_hz, time_mask_smooth_ms=time_mask_smooth_ms).to(device)
+
+    def _do_filter(self, chunk):
+        if type(chunk) is np.ndarray: chunk = torch.from_numpy(chunk).to(self.device)
+        return self.tg(x=chunk, xn=self.y_noise).cpu().detach().numpy()
+
+def reduce_noise(y, sr, stationary=False, y_noise=None, prop_decrease=1.0, time_constant_s=2.0, freq_mask_smooth_hz=500, time_mask_smooth_ms=50, thresh_n_mult_nonstationary=2, sigmoid_slope_nonstationary=10, tmp_folder=None, chunk_size=600000, padding=30000, n_fft=1024, win_length=None, hop_length=None, clip_noise_stationary=True, use_tqdm=False, device="cpu"):
+    return StreamedTorchGate(y=y, sr=sr, stationary=stationary, y_noise=y_noise, prop_decrease=prop_decrease, time_constant_s=time_constant_s, freq_mask_smooth_hz=freq_mask_smooth_hz, time_mask_smooth_ms=time_mask_smooth_ms, thresh_n_mult_nonstationary=thresh_n_mult_nonstationary, sigmoid_slope_nonstationary=sigmoid_slope_nonstationary, tmp_folder=tmp_folder, chunk_size=chunk_size, padding=padding, n_fft=n_fft, win_length=win_length, hop_length=hop_length, clip_noise_stationary=clip_noise_stationary, use_tqdm=use_tqdm, n_jobs=1, device=device).get_traces()