import ast import random from itertools import starmap, repeat from pathlib import Path from typing import Union import numpy as np import pandas as pd import soundfile as sf import torch import torchaudio import torchvision from hydra.utils import instantiate from omegaconf import DictConfig from torch.utils.data import Dataset from torchvision import transforms from audiomentations import Compose class BaseDataset(Dataset): """ class for storing and loading data. """ def __init__(self, data_path, metadata_path, augmentations, augmentations_p, preprocessors, label_type, seq_length=1, data_sample_rate=44100, sample_rate=44100, mode="train", slice_flag=False, margin_ratio=0, split_metadata_by_label=False, path_hierarchy: int = 0): """ __init__ method initiates ClassifierDataset instance: Input: data_path - string metadata_path - string augmentations - list of classes audiogemtations augmentations_p - array of probabilities (float64) preprocessors - list of classes from preprocessors (TBD function) path_hierarchy - enables working with data that is organized in a hierarchy of folders. The default value is 0, which means all the audio files are flattened in the same folder. If the value is 1, the audio files are organized in one folder per class, and so on. The annotations in the metadata has to be aligned with the path hierarchy, and to include the parent folder names in the filename column. Example: path_hierarchy = 0: - main_folder - file1.wav - file2.wav - file3.wav path_hierarchy = 1: - main_folder - sub_folder1 - file1.wav - file5.wav - sub_folder2 - file2.wav - file4.wav - sub_folder3 - file3.wav - file8.wav Output: ClassifierDataset Object - inherits from Dataset object in PyTorch package """ self.audio_dict = self._create_audio_dict(Path(data_path), path_hierarchy=path_hierarchy) self.metadata_path = metadata_path self.dtype_dict = {'filename': 'str'} self.label_type = label_type metadata = pd.read_csv(self.metadata_path, dtype=self.dtype_dict) self.metadata = self._update_metadata_by_mode(metadata, mode, split_metadata_by_label) self.mode = mode self.seq_length = seq_length self.sample_rate = sample_rate self.data_sample_rate = data_sample_rate self.sampler = torchaudio.transforms.Resample(orig_freq=data_sample_rate, new_freq=sample_rate) self._preprocess_metadata(slice_flag) self.augmenter = self._set_augmentations(augmentations, augmentations_p) self.preprocessor = self.set_preprocessor(preprocessors) assert (0 <= margin_ratio) and (1 >= margin_ratio) self.margin_ratio = margin_ratio self.num_classes = self._get_num_classes() self.samples_weight = self._get_samples_weight() @staticmethod def _update_metadata_by_mode(metadata, mode, split_metadata_by_label): if split_metadata_by_label: metadata = metadata[metadata['split_type'] == mode] return metadata def _create_audio_dict(self, data_path: Path, path_hierarchy=0) -> dict: """ create reference dict to extract audio files from metadata annotation Input: data_path - Path object Output: audio_dict contains references to audio paths given name from metadata """ def get_parent_path(path, path_hierarchy): parent_path_parts = path.parts[:-1] assert len(parent_path_parts) > path_hierarchy, \ (f"Make sure path_hierarchy:{path_hierarchy} is smaller than actual files hierarchy " f"{len(parent_path_parts)}") return '/'.join(parent_path_parts[len(parent_path_parts) - path_hierarchy:]) audio_paths = list(data_path.rglob('*.wav')) return {f'{get_parent_path(x, path_hierarchy)}/{x.name[:-4]}'.strip('/'): x for x in audio_paths} def _preprocess_metadata(self, slice_flag=False): """ function _preprocesses_metadata grabs calls with minimal length of self.seq_length + len_buffer Input: slice_flag: bool, default = False If true, the metadata file is sliced into segments of lengths self.seq_length. Output: ClassifierDataset object with self.metadata dataframe after applying the condition """ self.metadata['label'] = self._preprocess_target() is_noise = self.metadata['label'].apply(self._is_noise) # All calls are worthy (because we can later create a bigger slice contain them that is still a call in # _get_audio) but only long enough background sections will do. self.metadata = self.metadata[((self.metadata['call_length'] >= self.seq_length) & is_noise) | (~is_noise)] # sometimes the bbox's end time exceeds the file's length for name, sub_df in self.metadata.groupby('filename'): duration = sf.info(str(self.audio_dict[name])).duration if not all(sub_df['end_time'] <= duration): print(f'seems like some tags in file {name} have bigger end_time than its duration') print(f"file {name} --- int(duration): {int(duration)} --- biggest end time: {sub_df['end_time'].max()}") if slice_flag: self._slice_sequence() self.metadata.reset_index(drop=True, inplace=True) def _preprocess_target(self) -> pd.Series: """ Preprocesses the label column in the metadata. If the label is a string, it is evaluated and converted to an integer or a list of integers. """ if pd.api.types.is_string_dtype(self.metadata['label']): assert self.metadata['label'].str.match(r'^(\[|\()?(\d+)(\s*,\s*\d+)*(\]|\))?$').all(), \ "label should be a string that could be evaluated as a list of integers or integers." self.metadata['label'] = self.metadata['label'].apply(ast.literal_eval) if self.metadata['label'].apply(lambda x: isinstance(x, (list, tuple))).all(): self.metadata['label'] = self.metadata['label'].apply(np.array, dtype=int) return self.metadata['label'] @staticmethod def _is_noise(value: Union[int, np.ndarray]) -> bool: """ Checks if the value is a noise, i.e., if it is equal to 0. """ assert (isinstance(value, (int, np.integer)) | isinstance(value, np.ndarray)), "value should be either int or np.ndarray" return np.sum(value) == 0 def _grab_fields(self, idx): """ grabs fields from metadata according to idx input :idx output: begin_time - start time of segment end_time - end time of segment path_to_file - full path to file """ filename = self.metadata['filename'][idx] begin_time = self.metadata['begin_time'][idx] end_time = self.metadata['end_time'][idx] path_to_file = self.audio_dict[filename] orig_sample_rate = sf.info(path_to_file).samplerate assert orig_sample_rate == self.data_sample_rate begin_time = int(begin_time * orig_sample_rate) end_time = int(end_time * orig_sample_rate) label = self.metadata['label'][idx] if 'channel' in self.metadata.columns: channel = self.metadata['channel'][idx] else: channel = None return path_to_file, begin_time, end_time, label, channel def _slice_sequence(self): """ function _slice_sequence process metadata list call lengths to be sliced according to self.seq_length self Output: self.metadata sliced according to buffers """ self.metadata = self.metadata.reset_index(drop=True) count_values_before = self.metadata.astype({'label': str}).value_counts('label', sort=False) # for validating that the following code doesn't lose samples sliced_times = list(starmap(np.arange, zip(self.metadata['begin_time'], self.metadata['end_time'], repeat(self.seq_length)))) # add the last sequence at the end of this list for calls only (only if it does not exceed the file) sliced_times = list([np.append(s, self.metadata.loc[i, 'end_time']) if (not self._is_noise(self.metadata.loc[i, 'label'])) else s for i, s in enumerate(sliced_times)]) new_begin_time = list(x[:-1] for x in sliced_times) duplicate_size_vector = [len(list_elem) for list_elem in new_begin_time] # vector to duplicate original dataframe new_begin_time = np.concatenate(new_begin_time) # vectorize to array new_end_time = np.concatenate(list(x[1:] for x in sliced_times)) # same for end_times self.metadata = self.metadata.iloc[self.metadata.index.repeat(duplicate_size_vector)].reset_index(drop=True) self.metadata['begin_time'] = new_begin_time self.metadata['end_time'] = new_end_time self.metadata['call_length'] = np.shape(self.metadata)[0] * [self.seq_length] count_values_after = self.metadata.astype({'label': str}).value_counts('label', sort=False) if not all(count_values_after >= count_values_before): print(f'Note: seems like _slice_sequence erases data.\nbefore:{count_values_before}\n' f'after:{count_values_after}') return def _get_audio(self, path_to_file, begin_time, end_time, label, channel=None): raise NotImplementedError def _set_augmentations(self, augmentations_dict, augmentations_p): """ get augmentations list and instantiate - TBD """ if augmentations_dict is not None: augmentations_list = [instantiate(args) for args in augmentations_dict.values()] else: augmentations_list = [] self._train_augmenter = Compose(augmentations_list, p=augmentations_p, shuffle=True) self._val_augmenter = torch.nn.Identity() def augment(self, x): if self.mode == 'train': return torch.tensor(self._train_augmenter(x.numpy(), self.sample_rate), dtype=torch.float32) else: return self._val_augmenter(x) @staticmethod def set_preprocessor(preprocessors_args): """ function set_preprocessor takes preprocessors_args as an argument and creates a preprocessor object to be applied later on the audio segment input: preprocessors_args - list of classes from torchvision output: preprocessor - Composes several transforms together (transforms object) """ if len(preprocessors_args) > 0: processors_list = [instantiate(args) for args in preprocessors_args.values()] preprocessor = transforms.Compose(processors_list) else: preprocessor = torch.nn.Identity() return preprocessor def _get_num_classes(self) -> int: """ Returns the number of classes in the metadata. """ if self.label_type == 'multi_label': label_lengths = self.metadata['label'].apply(len) assert label_lengths.nunique() == 1, "All labels should have the same length" return label_lengths.iloc[0] else: return self.metadata['label'].nunique() def _get_samples_weight(self) -> np.ndarray: """ Returns the weight of each sample in the dataset: - if the label is integer, the weight is the inverse of the class count. - if the label is a list, the weight is the inverse of the minimum class count. """ if self.label_type == 'multi_label': noise_counts = self.metadata['label'].apply(self._is_noise).sum() class_counts = np.sum(self.metadata['label']) per_sample_min_class_count = (self.metadata['label'].apply( lambda x: class_counts[x.astype(bool)].min() if not self._is_noise(x) else noise_counts)) return (1 / per_sample_min_class_count).values else: weights = 1 / np.unique(self.metadata['label'], return_counts=True)[1] return np.array([weights[t] for t in self.metadata['label']]) def __getitem__(self, idx): ''' __getitem__ method loads item according to idx from the metadata input: idx - int output: For train/ val modes - audio_processed, label, audio_raw, idx - torch tensor (1-d if no spectrogram is applied/ 2-d if applied a spectrogram , int (if mode="train" only), 2-d tensor, int For test - audio_processed - torch tensor (1-d if no spectrogram is applied/ 2-d if applied a spectrogram ''' path_to_file, begin_time, end_time, label, channel = self._grab_fields(idx) audio = self._get_audio(path_to_file, begin_time, end_time, label, channel) audio_raw = self.sampler(audio) audio_augmented = self.augment(audio_raw) audio_processed = self.preprocessor(audio_augmented) if self.mode == "train" or self.mode == "val": label = self.metadata["label"][idx] return audio_processed, label, audio_raw, {"idx": idx, "begin_time": begin_time, "org_file": Path(path_to_file).stem} elif self.mode == "test": return audio_processed def __len__(self): return self.metadata.shape[0] class ClassifierDataset(BaseDataset): """ This class inherits all the traits from BaseDataset and handles cases that include Background noise (margin ratio feature is implemented) """ def _get_audio(self, path_to_file, begin_time, end_time, label, channel=None): """ _get_audio gets a path_to_file from _grab_fields method and also begin_time and end_time and returns the audio segment in a torch.tensor input: path_to_file - string begin_time - int end_time - int output: audio - pytorch tensor (1-D array) """ seg_length = end_time - begin_time requested_seq_length = int(self.seq_length * self.data_sample_rate) last_start_time = sf.info(path_to_file).frames - requested_seq_length # Do all this stuff only to calls in training set, because otherwise _slice_sequence has already been done if self.mode == "train": if seg_length >= requested_seq_length: # Only for calls we can safely add sections out of the call and label it as call if (self.margin_ratio != 0) and (not self._is_noise(label)): # self.margin_ratio ranges from 0 to 1 - indicates the relative part from seq_len to exceed call_length margin_len_begin = int(requested_seq_length * self.margin_ratio) margin_len_end = int(requested_seq_length * (1 - self.margin_ratio)) start_time = random.randint(max(begin_time - margin_len_begin, 0), min(end_time - margin_len_end, last_start_time)) else: start_time = random.randint(begin_time, min(end_time - requested_seq_length, last_start_time)) if start_time < 0: start_time = 0 else: # We know we can only arrive here with label > 0 because we filtered out short bg segments. # If the call is too short, the selected interval can be any interval of length requested_seq_length # that contains it. short_call_margin = requested_seq_length - seg_length # start time is between short_call_margin before begin time, and the latest time you can start and still # both contain the whole call and not get out of the file start_time = random.randint(max(begin_time - short_call_margin, 0), min(begin_time, last_start_time)) else: if begin_time < last_start_time: start_time = begin_time else: print(f'in {path_to_file}, one of the val\'s begin times is too big and exceeding the file so it was set to be smaller\nbegin time:{begin_time}, last_start_time:{last_start_time}') start_time = last_start_time data, _ = sf.read(str(path_to_file), start=start_time, stop=start_time + requested_seq_length) if channel is not None and data.ndim > 1: assert channel > 0, f"channel as to be a positive integer, got {channel}" data = data[:, channel - 1] elif channel is None and data.ndim > 1: data = data[:, 0] # when channel is not specified, take the first channel if data.shape[0] < 1: raise ValueError(f"Audio segment is empty. {path_to_file}: " f"{start_time}, {start_time + requested_seq_length}") audio = torch.tensor(data, dtype=torch.float).unsqueeze(0) return audio class NoBackGroundDataset(BaseDataset): """ This class inherits all the traits from BaseDataset and handles cases with no Background noise (calls only dataset) """ def _get_audio(self, path_to_file, begin_time, end_time, label, channel=None): """ _get_audio gets a path_to_file from _grab_fields method and also begin_time and end_time and returns the audio segment in a torch.tensor input: path_to_file - string begin_time - int end_time - int output: audio - pytorch tensor (1-D array) """ if (self.mode == "train"): start_time = random.randint(begin_time, end_time - int(self.seq_length * self.data_sample_rate)) if start_time < 0: start_time = 0 else: start_time = begin_time data, _ = sf.read(str(path_to_file), start=start_time, stop=start_time + int(self.seq_length * self.data_sample_rate)) if channel is not None: data = data[:, channel-1] if data.shape[0] < 1: raise ValueError(f"Audio segment is empty. {path_to_file}: " f"{start_time}, {start_time + int(self.seq_length * self.data_sample_rate)}") audio = torch.tensor(data, dtype=torch.float).unsqueeze(0) return audio class PeakNormalize: """Convert array to lay between 0 to 1""" def __call__(self, sample): return (sample - sample.min()) / (sample.max() - sample.min() + 1e-8) class MinFreqFiltering: """Cut the spectrogram frequency axis to make it start from min_freq ***Note: In case a MaxFreqFiltering is implemented, the max_freq should be greater than min_freq*** input: min_freq_filtering - int sample_rate - int output: spectrogram - pytorch tensor (3-D array) """ def __init__(self, min_freq_filtering, sample_rate): self.min_freq_filtering = min_freq_filtering self.sample_rate = sample_rate def edit_spectrogram_axis(self, sample): if self.min_freq_filtering > self.sample_rate / 2 or self.min_freq_filtering < 0: raise ValueError("min_freq_filtering should be greater than 0, and smaller than sample_rate/2") max_freq_in_spectrogram = self.sample_rate / 2 min_value = sample.size(dim=1) * self.min_freq_filtering / max_freq_in_spectrogram min_value = int(np.floor(min_value)) sample = sample[:, min_value:, :] return sample def __call__(self, sample): return self.edit_spectrogram_axis(sample) class UnitNormalize: """Remove mean and divide by std to normalize samples""" def __call__(self, sample): return (sample - sample.mean()) / (sample.std() + 1e-8) class SlidingWindowNormalize: """ Based on Sliding window augmentations of https://github.com/cchinchristopherj/Right-Whale-Convolutional-Neural-Network/blob/master/whale_cnn.py Translated to torch Has 50/50 chance of activating H sliding window or V sliding window Must come after spectrogram and before AmplitudeToDB """ def __init__(self, sr: float, n_fft: int, lower_cutoff: float = 50, norm=True, inner_ratio: float = 0.06, outer_ratio: float = 0.5): self.sr = sr self.n_fft = n_fft self.lower_cutoff = lower_cutoff self.norm = norm self.inner_ratio = inner_ratio self.outer_ratio = outer_ratio def spectrogram_norm(self, spect): min_f_ind = int((self.lower_cutoff / (self.sr / 2)) * self.n_fft) mval, sval = np.mean(spect[min_f_ind:, :]), np.std(spect[min_f_ind:, :]) fact_ = 1.5 spect[spect > mval + fact_ * sval] = mval + fact_ * sval spect[spect < mval - fact_ * sval] = mval - fact_ * sval spect[:min_f_ind, :] = mval return spect # slidingWindowV Function from: https://github.com/nmkridler/moby2/blob/master/metrics.py def slidingWindow(self, torch_spectrogram, dim=0): ''' slidingWindow Method Enhance the contrast vertically (along frequency dimension) for dim=0 and horizontally (along temporal dimension) for dim=1 Args: torch_spectrogram: 2-D numpy array image dim: dimension to do the sliding window across Returns: Q: 2-D numpy array image with vertically-enhanced contrast ''' if dim not in {0, 1}: raise ValueError('dim must be 0 or 1') spect = torch_spectrogram.cpu().clone().numpy() spect_shape = spect.shape spect = spect.squeeze() if self.norm: spect = self.spectrogram_norm(spect) # Set up the local mean window wInner = np.ones(int(self.inner_ratio * spect.shape[dim])) # Set up the overall mean window wOuter = np.ones(int(self.outer_ratio * spect.shape[dim])) # Remove overall mean and local mean using np.convolve for i in range(spect.shape[1-dim]): if dim == 0: spect[:, i] = spect[:, i] - ( np.convolve(spect[:, i], wOuter, 'same') - np.convolve(spect[:, i], wInner, 'same')) / ( wOuter.shape[0] - wInner.shape[0]) elif dim == 1: spect[i, :] = spect[i, :] - ( np.convolve(spect[i, :], wOuter, 'same') - np.convolve(spect[i, :], wInner, 'same')) / ( wOuter.shape[0] - wInner.shape[0]) spect[spect < 0] = 0. return torch.from_numpy(spect).reshape(spect_shape) def __call__(self, x): if random.random() < 0.5: return self.slidingWindow(x, dim=0) else: return self.slidingWindow(x, dim=1) class Resize(torchvision.transforms.Resize): def __init__(self, size): super().__init__(list(size)) class InferenceDataset(Dataset): ''' class for storing and loading data. ''' def __init__(self, file_path: Union[str, Path], preprocessors: DictConfig, seq_length: float = 1, data_sample_rate: int = 44100, sample_rate: int = 44100, overlap: float = 0): """ __init__ method initiates InferenceDataset instance: Input: Output: InferenceDataset Object - inherits from Dataset object in PyTorch package """ assert 0 <= overlap < 1, f'overlap should be between 0 and 1, got {overlap}' self.file_path = Path(file_path) self.metadata_path = self.file_path # alias to support inference pipeline self.seq_length = seq_length self.sample_rate = sample_rate self.data_sample_rate = data_sample_rate self.overlap = overlap self.sampler = torchaudio.transforms.Resample(orig_freq=data_sample_rate, new_freq=sample_rate) self.preprocessor = ClassifierDataset.set_preprocessor(preprocessors) self.metadata = self._create_inference_metadata() def _create_inference_metadata(self) -> pd.DataFrame: """ create metadata to be used in the inference dataset in case we have a directory, we will iterate over all files in the directory and create metadata for each file and merge it together For a single file, we will create metadata for that file """ all_data_frames = [] if self.file_path.is_dir(): all_files = [self.file_path / x for x in self.file_path.iterdir()] else: all_files = [self.file_path] for file in all_files: if file.suffix not in ['.wav', '.WAV']: raise ValueError(f'InferenceDataset only supports .wav files, got {file.suffix}') file_start_time = self._create_start_times(file) for channel_num in range(sf.info(file).channels): metadata = pd.DataFrame({'filename': [file] * len(file_start_time), 'channel': [channel_num] * len(file_start_time), 'begin_time': file_start_time, 'end_time': file_start_time + self.seq_length}) all_data_frames.append(metadata) metadata = pd.concat(all_data_frames, ignore_index=True) return metadata def _create_start_times(self, filepath: Path) -> np.ndarray: """ create reference dict to extract audio files from metadata annotation Input: data_path - Path object Output: audio_dict contains references to audio paths given name from metadata """ audio_len = sf.info(filepath).duration step = self.seq_length * (1-self.overlap) start_times = np.arange(0, audio_len, step) filtered_start_times = start_times[np.where(start_times <= audio_len - self.seq_length)] # if (duration - seq_length) is not a multiple of the step size, add the last segment if filtered_start_times[-1] < audio_len - self.seq_length: filtered_start_times = np.append(filtered_start_times, audio_len - self.seq_length) return filtered_start_times def _get_audio(self, filepath: Path, channel: int, begin_time: float) -> torch.Tensor: """ _get_audio gets a path_to_file from _grab_fields method and also begin_time and end_time and returns the audio segment in a torch.tensor input: path_to_file - string begin_time - int end_time - int output: audio - pytorch tensor (1-D array) """ duration = sf.info(filepath).duration begin_time = int(begin_time * self.data_sample_rate) stop_time = begin_time + int(self.seq_length * self.data_sample_rate) assert duration * self.data_sample_rate >= stop_time, f"trying to load audio from {begin_time} to {stop_time} but audio is only {duration} long" data, orig_sample_rate = sf.read(filepath, start=begin_time, stop=stop_time, always_2d=True) data = data[:, channel] assert orig_sample_rate == self.data_sample_rate, \ f'sample rate is {orig_sample_rate}, should be {self.data_sample_rate}' audio = torch.tensor(data, dtype=torch.float).unsqueeze(0) return audio def __getitem__(self, idx: int): ''' __getitem__ method loads item according to idx from the metadata. input: idx - int output: audio - torch tensor (1-d if no spectrogram is applied/ 2-d if applied a spectrogram ''' filepath, channel, begin_time = self.metadata.loc[idx, ['filename', 'channel', 'begin_time']] audio = self._get_audio(filepath=filepath, channel=channel, begin_time=begin_time) audio = self.sampler(audio) audio = self.preprocessor(audio) return audio def __len__(self): return len(self.metadata)