|
|
import torch |
|
|
import librosa |
|
|
import numpy as np |
|
|
|
|
|
class Denoiser: |
|
|
|
|
|
def __init__(self, orig_sr:int, target_sr:int): |
|
|
self.orig_sr = orig_sr |
|
|
self.target_sr = target_sr |
|
|
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
|
|
|
|
from asteroid.models import BaseModel as AsteroidBaseModel |
|
|
self.model = AsteroidBaseModel.from_pretrained("JorisCos/DCCRNet_Libri1Mix_enhsingle_16k").to(self.device) |
|
|
|
|
|
def denoise(self, wav): |
|
|
if type(wav) != np.ndarray: |
|
|
wav = np.array(wav) |
|
|
|
|
|
if len(wav.shape) > 1: |
|
|
wav = np.mean(wav, axis=1) |
|
|
wav = librosa.resample(wav, orig_sr=self.orig_sr, target_sr=self.target_sr) |
|
|
wav = torch.Tensor(wav.reshape(1, 1, wav.shape[0])).float().to(self.device) |
|
|
wav = self.model.separate(wav)[0][0] |
|
|
return wav.cpu().detach().numpy() |
|
|
|
