jone/Music_Source_Separation

1import sys2sys.path.append('.')3import argparse4import os5import time6from typing import Dict7import pathlib8 9import librosa10import numpy as np11import soundfile12import torch13import torch.nn as nn14 15from bytesep.models.lightning_modules import get_model_class16from bytesep.utils import read_yaml17 18 19class Separator:20    def __init__(21        self, model: nn.Module, segment_samples: int, batch_size: int, device: str22    ):23        r"""Separate to separate an audio clip into a target source.24 25        Args:26            model: nn.Module, trained model27            segment_samples: int, length of segments to be input to a model, e.g., 44100*3028            batch_size, int, e.g., 1229            device: str, e.g., 'cuda'30        """31        self.model = model32        self.segment_samples = segment_samples33        self.batch_size = batch_size34        self.device = device35 36    def separate(self, input_dict: Dict) -> np.array:37        r"""Separate an audio clip into a target source.38 39        Args:40            input_dict: dict, e.g., {41                waveform: (channels_num, audio_samples),42                ...,43            }44 45        Returns:46            sep_audio: (channels_num, audio_samples) | (target_sources_num, channels_num, audio_samples)47        """48        audio = input_dict['waveform']49 50        audio_samples = audio.shape[-1]51 52        # Pad the audio with zero in the end so that the length of audio can be53        # evenly divided by segment_samples.54        audio = self.pad_audio(audio)55 56        # Enframe long audio into segments.57        segments = self.enframe(audio, self.segment_samples)58        # (segments_num, channels_num, segment_samples)59 60        segments_input_dict = {'waveform': segments}61 62        if 'condition' in input_dict.keys():63            segments_num = len(segments)64            segments_input_dict['condition'] = np.tile(65                input_dict['condition'][None, :], (segments_num, 1)66            )67            # (batch_size, segments_num)68 69        # Separate in mini-batches.70        sep_segments = self._forward_in_mini_batches(71            self.model, segments_input_dict, self.batch_size72        )['waveform']73        # (segments_num, channels_num, segment_samples)74 75        # Deframe segments into long audio.76        sep_audio = self.deframe(sep_segments)77        # (channels_num, padded_audio_samples)78 79        sep_audio = sep_audio[:, 0:audio_samples]80        # (channels_num, audio_samples)81 82        return sep_audio83 84    def pad_audio(self, audio: np.array) -> np.array:85        r"""Pad the audio with zero in the end so that the length of audio can86        be evenly divided by segment_samples.87 88        Args:89            audio: (channels_num, audio_samples)90 91        Returns:92            padded_audio: (channels_num, audio_samples)93        """94        channels_num, audio_samples = audio.shape95 96        # Number of segments97        segments_num = int(np.ceil(audio_samples / self.segment_samples))98 99        pad_samples = segments_num * self.segment_samples - audio_samples100 101        padded_audio = np.concatenate(102            (audio, np.zeros((channels_num, pad_samples))), axis=1103        )104        # (channels_num, padded_audio_samples)105 106        return padded_audio107 108    def enframe(self, audio: np.array, segment_samples: int) -> np.array:109        r"""Enframe long audio into segments.110 111        Args:112            audio: (channels_num, audio_samples)113            segment_samples: int114 115        Returns:116            segments: (segments_num, channels_num, segment_samples)117        """118        audio_samples = audio.shape[1]119        assert audio_samples % segment_samples == 0120 121        hop_samples = segment_samples // 2122        segments = []123 124        pointer = 0125        while pointer + segment_samples <= audio_samples:126            segments.append(audio[:, pointer : pointer + segment_samples])127            pointer += hop_samples128 129        segments = np.array(segments)130 131        return segments132 133    def deframe(self, segments: np.array) -> np.array:134        r"""Deframe segments into long audio.135 136        Args:137            segments: (segments_num, channels_num, segment_samples)138 139        Returns:140            output: (channels_num, audio_samples)141        """142        (segments_num, _, segment_samples) = segments.shape143 144        if segments_num == 1:145            return segments[0]146 147        assert self._is_integer(segment_samples * 0.25)148        assert self._is_integer(segment_samples * 0.75)149 150        output = []151 152        output.append(segments[0, :, 0 : int(segment_samples * 0.75)])153 154        for i in range(1, segments_num - 1):155            output.append(156                segments[157                    i, :, int(segment_samples * 0.25) : int(segment_samples * 0.75)158                ]159            )160 161        output.append(segments[-1, :, int(segment_samples * 0.25) :])162 163        output = np.concatenate(output, axis=-1)164 165        return output166 167    def _is_integer(self, x: float) -> bool:168        if x - int(x) < 1e-10:169            return True170        else:171            return False172 173    def _forward_in_mini_batches(174        self, model: nn.Module, segments_input_dict: Dict, batch_size: int175    ) -> Dict:176        r"""Forward data to model in mini-batch.177 178        Args:179            model: nn.Module180            segments_input_dict: dict, e.g., {181                'waveform': (segments_num, channels_num, segment_samples),182                ...,183            }184            batch_size: int185 186        Returns:187            output_dict: dict, e.g. {188                'waveform': (segments_num, channels_num, segment_samples),189            }190        """191        output_dict = {}192 193        pointer = 0194        segments_num = len(segments_input_dict['waveform'])195 196        while True:197            if pointer >= segments_num:198                break199 200            batch_input_dict = {}201 202            for key in segments_input_dict.keys():203                batch_input_dict[key] = torch.Tensor(204                    segments_input_dict[key][pointer : pointer + batch_size]205                ).to(self.device)206 207            pointer += batch_size208 209            with torch.no_grad():210                model.eval()211                batch_output_dict = model(batch_input_dict)212 213            for key in batch_output_dict.keys():214                self._append_to_dict(215                    output_dict, key, batch_output_dict[key].data.cpu().numpy()216                )217 218        for key in output_dict.keys():219            output_dict[key] = np.concatenate(output_dict[key], axis=0)220 221        return output_dict222 223    def _append_to_dict(self, dict, key, value):224        if key in dict.keys():225            dict[key].append(value)226        else:227            dict[key] = [value]228 229 230class SeparatorWrapper:231    def __init__(232        self, source_type='vocals', model=None, checkpoint_path=None, device='cuda'233    ):234 235        input_channels = 2236        target_sources_num = 1237        model_type = "ResUNet143_Subbandtime"238        segment_samples = 44100 * 10239        batch_size = 1240 241        self.checkpoint_path = self.download_checkpoints(checkpoint_path, source_type)242 243        if device == 'cuda' and torch.cuda.is_available():244            self.device = 'cuda'245        else:246            self.device = 'cpu'247 248        # Get model class.249        Model = get_model_class(model_type)250 251        # Create model.252        self.model = Model(253            input_channels=input_channels, target_sources_num=target_sources_num254        )255 256        # Load checkpoint.257        checkpoint = torch.load(self.checkpoint_path, map_location='cpu')258        self.model.load_state_dict(checkpoint["model"])259 260        # Move model to device.261        self.model.to(self.device)262 263        # Create separator.264        self.separator = Separator(265            model=self.model,266            segment_samples=segment_samples,267            batch_size=batch_size,268            device=self.device,269        )270 271    def download_checkpoints(self, checkpoint_path, source_type):272 273        if source_type == "vocals":274            checkpoint_bare_name = "resunet143_subbtandtime_vocals_8.8dB_350k_steps"275 276        elif source_type == "accompaniment":277            checkpoint_bare_name = (278                "resunet143_subbtandtime_accompaniment_16.4dB_350k_steps.pth"279            )280 281        else:282            raise NotImplementedError283 284        if not checkpoint_path:285            checkpoint_path = '{}/bytesep_data/{}.pth'.format(286                str(pathlib.Path.home()), checkpoint_bare_name287            )288 289        print('Checkpoint path: {}'.format(checkpoint_path))290 291        if (292            not os.path.exists(checkpoint_path)293            or os.path.getsize(checkpoint_path) < 4e8294        ):295 296            os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True)297 298            zenodo_dir = "https://zenodo.org/record/5507029/files"299            zenodo_path = os.path.join(300                zenodo_dir, "{}?download=1".format(checkpoint_bare_name)301            )302 303            os.system('wget -O "{}" "{}"'.format(checkpoint_path, zenodo_path))304 305        return checkpoint_path306 307    def separate(self, audio):308 309        input_dict = {'waveform': audio}310 311        sep_wav = self.separator.separate(input_dict)312 313        return sep_wav314 315 316def inference(args):317 318    # Need to use torch.distributed if models contain inplace_abn.abn.InPlaceABNSync.319    import torch.distributed as dist320 321    dist.init_process_group(322        'gloo', init_method='file:///tmp/somefile', rank=0, world_size=1323    )324 325    # Arguments & parameters326    config_yaml = args.config_yaml327    checkpoint_path = args.checkpoint_path328    audio_path = args.audio_path329    output_path = args.output_path330    device = (331        torch.device('cuda')332        if args.cuda and torch.cuda.is_available()333        else torch.device('cpu')334    )335 336    configs = read_yaml(config_yaml)337    sample_rate = configs['train']['sample_rate']338    input_channels = configs['train']['channels']339    target_source_types = configs['train']['target_source_types']340    target_sources_num = len(target_source_types)341    model_type = configs['train']['model_type']342 343    segment_samples = int(30 * sample_rate)344    batch_size = 1345 346    print("Using {} for separating ..".format(device))347 348    # paths349    if os.path.dirname(output_path) != "":350        os.makedirs(os.path.dirname(output_path), exist_ok=True)351 352    # Get model class.353    Model = get_model_class(model_type)354 355    # Create model.356    model = Model(input_channels=input_channels, target_sources_num=target_sources_num)357 358    # Load checkpoint.359    checkpoint = torch.load(checkpoint_path, map_location='cpu')360    model.load_state_dict(checkpoint["model"])361 362    # Move model to device.363    model.to(device)364 365    # Create separator.366    separator = Separator(367        model=model,368        segment_samples=segment_samples,369        batch_size=batch_size,370        device=device,371    )372 373    # Load audio.374    audio, _ = librosa.load(audio_path, sr=sample_rate, mono=False)375 376    # audio = audio[None, :]377 378    input_dict = {'waveform': audio}379 380    # Separate381    separate_time = time.time()382 383    sep_wav = separator.separate(input_dict)384    # (channels_num, audio_samples)385 386    print('Separate time: {:.3f} s'.format(time.time() - separate_time))387 388    # Write out separated audio.389    soundfile.write(file='_zz.wav', data=sep_wav.T, samplerate=sample_rate)390    os.system("ffmpeg -y -loglevel panic -i _zz.wav {}".format(output_path))391    print('Write out to {}'.format(output_path))392 393 394if __name__ == "__main__":395 396    parser = argparse.ArgumentParser(description="")397    parser.add_argument("--config_yaml", type=str, required=True)398    parser.add_argument("--checkpoint_path", type=str, required=True)399    parser.add_argument("--audio_path", type=str, required=True)400    parser.add_argument("--output_path", type=str, required=True)401    parser.add_argument("--cuda", action='store_true', default=True)402 403    args = parser.parse_args()404    inference(args)405