jone/Music_Source_Separation

1from typing import Dict, List2 3import torch4 5 6class BasicBatchDataPreprocessor:7    def __init__(self, target_source_types: List[str]):8        r"""Batch data preprocessor. Used for preparing mixtures and targets for9        training. If there are multiple target source types, the waveforms of10        those sources will be stacked along the channel dimension.11 12        Args:13            target_source_types: List[str], e.g., ['vocals', 'bass', ...]14        """15        self.target_source_types = target_source_types16 17    def __call__(self, batch_data_dict: Dict) -> List[Dict]:18        r"""Format waveforms and targets for training.19 20        Args:21            batch_data_dict: dict, e.g., {22                'mixture': (batch_size, channels_num, segment_samples),23                'vocals': (batch_size, channels_num, segment_samples),24                'bass': (batch_size, channels_num, segment_samples),25                ...,26            }27 28        Returns:29            input_dict: dict, e.g., {30                'waveform': (batch_size, channels_num, segment_samples),31            }32            output_dict: dict, e.g., {33                'target': (batch_size, target_sources_num * channels_num, segment_samples)34            }35        """36        mixtures = batch_data_dict['mixture']37        # mixtures: (batch_size, channels_num, segment_samples)38 39        # Concatenate waveforms of multiple targets along the channel axis.40        targets = torch.cat(41            [batch_data_dict[source_type] for source_type in self.target_source_types],42            dim=1,43        )44        # targets: (batch_size, target_sources_num * channels_num, segment_samples)45 46        input_dict = {'waveform': mixtures}47        target_dict = {'waveform': targets}48 49        return input_dict, target_dict50 51 52class ConditionalSisoBatchDataPreprocessor:53    def __init__(self, target_source_types: List[str]):54        r"""Conditional single input single output (SISO) batch data55        preprocessor. Select one target source from several target sources as56        training target and prepare the corresponding conditional vector.57 58        Args:59            target_source_types: List[str], e.g., ['vocals', 'bass', ...]60        """61        self.target_source_types = target_source_types62 63    def __call__(self, batch_data_dict: Dict) -> List[Dict]:64        r"""Format waveforms and targets for training.65 66        Args:67            batch_data_dict: dict, e.g., {68                'mixture': (batch_size, channels_num, segment_samples),69                'vocals': (batch_size, channels_num, segment_samples),70                'bass': (batch_size, channels_num, segment_samples),71                ...,72            }73 74        Returns:75            input_dict: dict, e.g., {76                'waveform': (batch_size, channels_num, segment_samples),77                'condition': (batch_size, target_sources_num),78            }79            output_dict: dict, e.g., {80                'target': (batch_size, channels_num, segment_samples)81            }82        """83 84        batch_size = len(batch_data_dict['mixture'])85        target_sources_num = len(self.target_source_types)86 87        assert (88            batch_size % target_sources_num == 089        ), "Batch size should be \90            evenly divided by target sources number."91 92        mixtures = batch_data_dict['mixture']93        # mixtures: (batch_size, channels_num, segment_samples)94 95        conditions = torch.zeros(batch_size, target_sources_num).to(mixtures.device)96        # conditions: (batch_size, target_sources_num)97 98        targets = []99 100        for n in range(batch_size):101 102            k = n % target_sources_num  # source class index103            source_type = self.target_source_types[k]104 105            targets.append(batch_data_dict[source_type][n])106 107            conditions[n, k] = 1108 109        # conditions will looks like:110        # [[1, 0, 0, 0],111        #  [0, 1, 0, 0],112        #  [0, 0, 1, 0],113        #  [0, 0, 0, 1],114        #  [1, 0, 0, 0],115        #  [0, 1, 0, 0],116        #  ...,117        # ]118 119        targets = torch.stack(targets, dim=0)120        # targets: (batch_size, channels_num, segment_samples)121 122        input_dict = {123            'waveform': mixtures,124            'condition': conditions,125        }126 127        target_dict = {'waveform': targets}128 129        return input_dict, target_dict130 131 132def get_batch_data_preprocessor_class(batch_data_preprocessor_type: str) -> object:133    r"""Get batch data preprocessor class."""134    if batch_data_preprocessor_type == 'BasicBatchDataPreprocessor':135        return BasicBatchDataPreprocessor136 137    elif batch_data_preprocessor_type == 'ConditionalSisoBatchDataPreprocessor':138        return ConditionalSisoBatchDataPreprocessor139 140    else:141        raise NotImplementedError142