ASesYusuf1/SESA_Audio_Separation

1import torch2import torch.nn as nn3import segmentation_models_pytorch as smp4from utils import prefer_target_instrument5 6class STFT:7    def __init__(self, config):8        self.n_fft = config.n_fft9        self.hop_length = config.hop_length10        self.window = torch.hann_window(window_length=self.n_fft, periodic=True)11        self.dim_f = config.dim_f12 13    def __call__(self, x):14        window = self.window.to(x.device)15        batch_dims = x.shape[:-2]16        c, t = x.shape[-2:]17        x = x.reshape([-1, t])18        x = torch.stft(19            x,20            n_fft=self.n_fft,21            hop_length=self.hop_length,22            window=window,23            center=True,24            return_complex=True25        )26        x = torch.view_as_real(x)27        x = x.permute([0, 3, 1, 2])28        x = x.reshape([*batch_dims, c, 2, -1, x.shape[-1]]).reshape([*batch_dims, c * 2, -1, x.shape[-1]])29        return x[..., :self.dim_f, :]30 31    def inverse(self, x):32        window = self.window.to(x.device)33        batch_dims = x.shape[:-3]34        c, f, t = x.shape[-3:]35        n = self.n_fft // 2 + 136        f_pad = torch.zeros([*batch_dims, c, n - f, t]).to(x.device)37        x = torch.cat([x, f_pad], -2)38        x = x.reshape([*batch_dims, c // 2, 2, n, t]).reshape([-1, 2, n, t])39        x = x.permute([0, 2, 3, 1])40        x = x[..., 0] + x[..., 1] * 1.j41        x = torch.istft(42            x,43            n_fft=self.n_fft,44            hop_length=self.hop_length,45            window=window,46            center=True47        )48        x = x.reshape([*batch_dims, 2, -1])49        return x50 51 52def get_act(act_type):53    if act_type == 'gelu':54        return nn.GELU()55    elif act_type == 'relu':56        return nn.ReLU()57    elif act_type[:3] == 'elu':58        alpha = float(act_type.replace('elu', ''))59        return nn.ELU(alpha)60    else:61        raise Exception62 63 64def get_decoder(config, c):65    decoder = None66    decoder_options = dict()67    if config.model.decoder_type == 'unet':68        try:69            decoder_options = dict(config.decoder_unet)70        except:71            pass72        decoder = smp.Unet(73            encoder_name=config.model.encoder_name,74            encoder_weights="imagenet",75            in_channels=c,76            classes=c,77            **decoder_options,78        )79    elif config.model.decoder_type == 'fpn':80        try:81            decoder_options = dict(config.decoder_fpn)82        except:83            pass84        decoder = smp.FPN(85            encoder_name=config.model.encoder_name,86            encoder_weights="imagenet",87            in_channels=c,88            classes=c,89            **decoder_options,90        )91    elif config.model.decoder_type == 'unet++':92        try:93            decoder_options = dict(config.decoder_unet_plus_plus)94        except:95            pass96        decoder = smp.UnetPlusPlus(97            encoder_name=config.model.encoder_name,98            encoder_weights="imagenet",99            in_channels=c,100            classes=c,101            **decoder_options,102        )103    elif config.model.decoder_type == 'manet':104        try:105            decoder_options = dict(config.decoder_manet)106        except:107            pass108        decoder = smp.MAnet(109            encoder_name=config.model.encoder_name,110            encoder_weights="imagenet",111            in_channels=c,112            classes=c,113            **decoder_options,114        )115    elif config.model.decoder_type == 'linknet':116        try:117            decoder_options = dict(config.decoder_linknet)118        except:119            pass120        decoder = smp.Linknet(121            encoder_name=config.model.encoder_name,122            encoder_weights="imagenet",123            in_channels=c,124            classes=c,125            **decoder_options,126        )127    elif config.model.decoder_type == 'pspnet':128        try:129            decoder_options = dict(config.decoder_pspnet)130        except:131            pass132        decoder = smp.PSPNet(133            encoder_name=config.model.encoder_name,134            encoder_weights="imagenet",135            in_channels=c,136            classes=c,137            **decoder_options,138        )139    elif config.model.decoder_type == 'pspnet':140        try:141            decoder_options = dict(config.decoder_pspnet)142        except:143            pass144        decoder = smp.PSPNet(145            encoder_name=config.model.encoder_name,146            encoder_weights="imagenet",147            in_channels=c,148            classes=c,149            **decoder_options,150        )151    elif config.model.decoder_type == 'pan':152        try:153            decoder_options = dict(config.decoder_pan)154        except:155            pass156        decoder = smp.PAN(157            encoder_name=config.model.encoder_name,158            encoder_weights="imagenet",159            in_channels=c,160            classes=c,161            **decoder_options,162        )163    elif config.model.decoder_type == 'deeplabv3':164        try:165            decoder_options = dict(config.decoder_deeplabv3)166        except:167            pass168        decoder = smp.DeepLabV3(169            encoder_name=config.model.encoder_name,170            encoder_weights="imagenet",171            in_channels=c,172            classes=c,173            **decoder_options,174        )175    elif config.model.decoder_type == 'deeplabv3plus':176        try:177            decoder_options = dict(config.decoder_deeplabv3plus)178        except:179            pass180        decoder = smp.DeepLabV3Plus(181            encoder_name=config.model.encoder_name,182            encoder_weights="imagenet",183            in_channels=c,184            classes=c,185            **decoder_options,186        )187    return decoder188 189 190class Segm_Models_Net(nn.Module):191    def __init__(self, config):192        super().__init__()193        self.config = config194 195        act = get_act(act_type=config.model.act)196 197        self.num_target_instruments = len(prefer_target_instrument(config))198        self.num_subbands = config.model.num_subbands199 200        dim_c = self.num_subbands * config.audio.num_channels * 2201        c = config.model.num_channels202        f = config.audio.dim_f // self.num_subbands203 204        self.first_conv = nn.Conv2d(dim_c, c, 1, 1, 0, bias=False)205 206        self.unet_model = get_decoder(config, c)207 208        self.final_conv = nn.Sequential(209            nn.Conv2d(c + dim_c, c, 1, 1, 0, bias=False),210            act,211            nn.Conv2d(c, self.num_target_instruments * dim_c, 1, 1, 0, bias=False)212        )213 214        self.stft = STFT(config.audio)215 216    def cac2cws(self, x):217        k = self.num_subbands218        b, c, f, t = x.shape219        x = x.reshape(b, c, k, f // k, t)220        x = x.reshape(b, c * k, f // k, t)221        return x222 223    def cws2cac(self, x):224        k = self.num_subbands225        b, c, f, t = x.shape226        x = x.reshape(b, c // k, k, f, t)227        x = x.reshape(b, c // k, f * k, t)228        return x229 230    def forward(self, x):231 232        x = self.stft(x)233 234        mix = x = self.cac2cws(x)235 236        first_conv_out = x = self.first_conv(x)237 238        x = x.transpose(-1, -2)239 240        x = self.unet_model(x)241 242        x = x.transpose(-1, -2)243 244        x = x * first_conv_out  # reduce artifacts245 246        x = self.final_conv(torch.cat([mix, x], 1))247 248        x = self.cws2cac(x)249 250        if self.num_target_instruments > 1:251            b, c, f, t = x.shape252            x = x.reshape(b, self.num_target_instruments, -1, f, t)253 254        x = self.stft.inverse(x)255        return x256