Adding the dataloader. Has not been tested in this context, TODO

as_seg/model/dataloaders.py

+import mirdata
+import librosa
+import as_seg.model.signal_to_spectrogram as signal_to_spectrogram
+import pathlib
+import shutil
+import numpy as np
+import warnings
+
+import as_seg
+
+eps = 1e-10
+
+class BaseDataloader():
+    def __init__(self, feature, cache_path = None, sr=44100, hop_length = 32, subdivision = 96, verbose = False):
+        self.cache_path = cache_path
+        self.verbose = verbose
+
+        self.sr = sr
+        self.feature = feature
+        self.hop_length = hop_length
+
+        # For barwise or beatwise processing
+        self.subdivision = subdivision
+
+        self.frequency_dimension = signal_to_spectrogram.get_default_frequency_dimension(feature) # Risky, because it is not linked to the computation. Should be computed from the spectrogram.
+
+    def __getitem__(self, index):
+        raise NotImplementedError("This method should be implemented in the child class") from None
+
+    def __len__(self):
+        raise NotImplementedError("This method should be implemented in the child class") from None
+
+    def get_spectrogram(self, signal): # The spectrogram is not saved in the cache because it is too large in general
+        return signal_to_spectrogram.get_spectrogram(signal, self.sr, self.feature, self.hop_length)
+    
+    def get_bars(self, audio_path, index = None):
+        def _compute_bars():
+            return as_seg.data_manipulation.get_bars_from_audio(audio_path)
+
+        if self.cache_path is not None:
+            if index is None:
+                warnings.warn("No index provided for the cache, the cache will be ignored")
+            else:
+                dir_save_bars_path = f"{self.cache_path}/bars"
+                try:
+                    bars = np.load(f"{dir_save_bars_path}/{index}.npy", allow_pickle=True)
+                    if self.verbose:
+                        print("Using cached bars.")
+                except FileNotFoundError:
+                    bars = _compute_bars()
+                    pathlib.Path(dir_save_bars_path).mkdir(parents=True, exist_ok=True)
+                    np.save(f"{dir_save_bars_path}/{index}.npy", bars)
+                return bars
+        return _compute_bars()
+
+    def get_barwise_tf_matrix(self, track_path, bars, index = None):
+        def _compute_barwise_tf_matrix():
+            # Load the signal of the song
+            sig, _ = librosa.load(track_path, sr=self.sr, mono=True) #torchaudio.load(track.audio_path)
+            # Compute the spectrogram
+            spectrogram = self.get_spectrogram(sig)
+            return as_seg.barwise_input.barwise_TF_matrix(spectrogram, bars, self.hop_length/self.sr, self.subdivision) + eps
+        
+        if self.cache_path is not None:
+            if index is None:
+                warnings.warn("No index provided for the cache, the cache will be ignored")
+            else:
+                cache_file_name = f"{index}_{self.feature}_subdiv{self.subdivision}"
+                dir_save_barwise_tf_path = f"{self.cache_path}/barwise_tf_matrix"
+                try:
+                    barwise_tf_matrix = np.load(f"{dir_save_barwise_tf_path}/{cache_file_name}.npy", allow_pickle=True)
+                    if self.verbose:
+                        print("Using cached Barwise TF matrix.")
+                except FileNotFoundError:
+                    barwise_tf_matrix = _compute_barwise_tf_matrix()
+                    pathlib.Path(dir_save_barwise_tf_path).mkdir(parents=True, exist_ok=True)
+                    np.save(f"{dir_save_barwise_tf_path}/{cache_file_name}.npy", barwise_tf_matrix)
+                return barwise_tf_matrix
+        return _compute_barwise_tf_matrix()
+
+    def save_segments(self, segments, name):
+        # mirdata_segments = mirdata.annotations.SectionData(intervals=segments, interval_unit="s")
+        # jams_segments = mirdata.jams_utils.sections_to_jams(mirdata_segments)
+        dir_save_path = f"{self.data_path}/estimations/segments/{self.dataset_name.lower()}"
+        pathlib.Path(dir_save_path).mkdir(parents=True, exist_ok=True)
+        np.save(f"{dir_save_path}/{name}.npy", segments)
+
+    def score_flat_segmentation(self, segments, annotations):
+        close_tolerance = as_seg.data_manipulation.compute_score_of_segmentation(annotations, segments, window_length=0.5)
+        large_tolerance = as_seg.data_manipulation.compute_score_of_segmentation(annotations, segments, window_length=3)
+        return close_tolerance, large_tolerance
+    
+    def segments_from_bar_to_seconds(self, segments, bars):
+        # May be useful, if ever.
+        return as_seg.data_manipulation.segments_from_bar_to_time(segments, bars)
+
+class RWCPopDataloader(BaseDataloader):
+    def __init__(self, path, feature, cache_path = None, download=False, sr=44100, hop_length = 32, subdivision = 96):
+        super().__init__(feature, cache_path, sr, hop_length, subdivision)
+        self.data_path = path
+        rwcpop = mirdata.initialize('rwc_popular', data_home = path)
+        if download:
+            rwcpop.download()            
+        self.all_tracks = rwcpop.load_tracks()
+        self.indexes = rwcpop.track_ids
+
+        self.dataset_name = "RWCPop"
+
+    def __getitem__(self, index):
+        track_id = self.indexes[index]
+        track = self.all_tracks[track_id]
+
+        # Compute the bars
+        bars = self.get_bars(track.audio_path, index=track_id)
+
+        # Compute the barwise TF matrix
+        barwise_tf_matrix = self.get_barwise_tf_matrix(track.audio_path, bars, index=track_id)
+
+        # Get the annotationks
+        annotations_intervals = track.sections.intervals
+
+        # Return the the bars, the barwise TF matrix and the annotations
+        return track_id, bars, barwise_tf_matrix, annotations_intervals
+    
+    def __len__(self):
+        return len(self.indexes)
+    
+    def get_track_of_id(self, track_id):
+        index = self.indexes.index(track_id)
+        return self.__getitem__[index]
+    
+    def format_dataset(self, path_audio_files):
+        # Copy audio files to the right location.
+        # Suppose that the audio files are all in the same folder
+        for track_num in range(len(self.all_tracks)):
+            track_idx = self.indexes[track_num]
+            song_file_name = self.all_tracks[track_idx].audio_path.split('/')[-1]
+            src = f"{path_audio_files}/{song_file_name}" # May change depending on your file structure
+            dest = self.all_tracks[track_idx].audio_path
+            pathlib.Path(dest).parent.absolute().mkdir(parents=True, exist_ok=True)
+            shutil.copy(src, dest)
+
+class SALAMIDataloader(BaseDataloader):
+    def __init__(self, path, feature, cache_path = None, download=False, subset = None, sr=44100, hop_length = 32, subdivision = 96):
+        super().__init__(feature, cache_path, sr, hop_length, subdivision)
+        
+        self.dataset_name = "SALAMI"
+
+        self.data_path = path
+        salami = mirdata.initialize('salami', data_home = path)
+        if download:
+            salami.download()            
+        self.all_tracks = salami.load_tracks()
+        self.indexes = salami.track_ids
+
+        self.subset = subset
+        if subset is not None:
+            train_indexes, test_indexes = self.split_training_test()
+            if subset == "train":
+                self.indexes = train_indexes
+            elif subset == "test":
+                self.indexes = test_indexes
+            elif subset == "debug":
+                self.indexes = test_indexes[:4]
+            else:
+                raise ValueError("Subset should be either 'train' or 'test'")
+
+
+    def __getitem__(self, index):
+        # Parsing through files ordered with self.indexes
+        track_id = self.indexes[index]
+        track = self.all_tracks[track_id]
+
+        try:           
+            # Compute the bars
+            bars = self.get_bars(track.audio_path, index=track_id)
+
+            # Compute the barwise TF matrix
+            barwise_tf_matrix = self.get_barwise_tf_matrix(track.audio_path, bars, index=track_id)
+
+            # Get the annotations
+            dict_annotations = self.get_annotations(track)
+
+            # Return the the bars, the barwise TF matrix and the annotations
+            return track_id, bars, barwise_tf_matrix, dict_annotations
+    
+        except FileNotFoundError:
+            print(f'{track_id} not found.')
+            return track_id, None, None, None
+            # raise FileNotFoundError(f"Song {track_id} not found, normal ?") from None
+    
+    def __len__(self):
+        # To handle the fact that indexes are updated with the subset
+        return len(self.indexes)
+
+    def get_track_of_id(self, track_id):
+        try:
+            index = self.indexes.index(track_id)
+        except ValueError:
+            try:
+                index = self.indexes.index(str(track_id))
+            except ValueError:
+                raise ValueError(f"Track {track_id} not found in the dataset") from None
+        return self.__getitem__(index)
+
+    def get_annotations(self, track):
+        dict_annotations = {}
+        try: 
+            # Trying to get the first annotator
+            dict_annotations["upper_level_annotations"] = track.sections_annotator_1_uppercase.intervals
+            dict_annotations["lower_level_annotations"] = track.sections_annotator_1_lowercase.intervals
+            try: # Trying to load the second annotator
+                dict_annotations["upper_level_annotations_2"] = track.sections_annotator_2_uppercase.intervals
+                dict_annotations["lower_level_annotations_2"] = track.sections_annotator_2_lowercase.intervals
+                dict_annotations["annot_number"]  = 2
+            except AttributeError: # Only the first annotator was loaded
+                dict_annotations["annot_number"]  = 1
+        except AttributeError:
+            try:
+                # Trying to get the second annotator (no first one)
+                dict_annotations["upper_level_annotations"] = track.sections_annotator_2_uppercase.intervals
+                dict_annotations["lower_level_annotations"] = track.sections_annotator_2_lowercase.intervals
+                dict_annotations["annot_number"]  = 1
+            except AttributeError:
+                raise AttributeError(f"No annotations found for {track.track_id}")
+        
+        return dict_annotations
+    
+    def get_this_set_annotations(self, dict_annotations, annotation_level = "upper", annotator = 1):
+        if annotator == 1:
+            if annotation_level == "upper":
+                annotations = dict_annotations["upper_level_annotations"]
+            elif annotation_level == "lower":
+                annotations = dict_annotations["lower_level_annotations"]
+            else:
+                raise ValueError("Invalid annotation level")
+        elif annotator == 2:
+            assert dict_annotations["annot_number"] == 2, "No second annotator found."
+            if annotation_level == "upper":
+                annotations = dict_annotations["upper_level_annotations"]
+            elif annotation_level == "lower":
+                annotations = dict_annotations["lower_level_annotations"]
+            else:
+                raise ValueError("Invalid annotation level")
+        # elif annotator == "both":
+        #     assert dict_annotations["annot_number"] == 2, "No second annotator found."
+        #     annotations = dict_annotations["upper_level_annotations"] + dict_annotations["upper_level_annotations_2"]
+        else:
+            raise ValueError("Invalid annotator number")
+        return annotations
+
+    def split_training_test(self):
+        indexes_train = []
+        indexes_test = []
+        for track_id in self.indexes:
+            track = self.all_tracks[track_id]
+            try:
+                track.sections_annotator_1_uppercase.intervals
+                track.sections_annotator_2_uppercase.intervals
+                indexes_test.append(track_id)
+            except AttributeError:
+                indexes_train.append(track_id)
+        return indexes_train, indexes_test
+    
+    def score_flat_segmentation(self, segments, dict_annotations, annotation_level = "upper", annotator = 1):
+        if annotator == "both":
+            assert dict_annotations["annot_number"] == 2, "No second annotator found."
+            score_annot_1 = self.score_flat_segmentation(segments, dict_annotations, annotation_level = annotation_level, annotator = 1)
+            score_annot_2 = self.score_flat_segmentation(segments, dict_annotations, annotation_level = annotation_level, annotator = 2)
+            return score_annot_1, score_annot_2
+        
+        annotations = self.get_this_set_annotations(dict_annotations, annotation_level = annotation_level, annotator = annotator)
+        return super().score_flat_segmentation(segments, annotations)
+        
+    def score_flat_segmentation_twolevel(self, segments_upper_level, segments_lower_level, dict_annotations, annotator = 1):
+        score_upper_level = self.score_flat_segmentation(segments_upper_level, dict_annotations, annotation_level = "upper", annotator = annotator)
+        score_lower_level = self.score_flat_segmentation(segments_lower_level, dict_annotations, annotation_level = "lower", annotator = annotator)
+        return score_upper_level, score_lower_level
+    
+    def score_flat_segmentation_twolevel_best_of_several(self, list_segments_upper_level, list_segments_lower_level, dict_annotations, annotator = 1):
+        assert annotator != "both", "Not implemented yet"
+        stack_upper_scores = -np.inf * np.ones((len(list_segments_upper_level),2,3))
+        for idx, segments in enumerate(list_segments_upper_level):
+            stack_upper_scores[idx] = self.score_flat_segmentation(segments, dict_annotations, annotation_level = "upper", annotator = annotator)
+        idx_close = np.argmax(stack_upper_scores[:,0,2]) # Selecting best f measure at 0.5s
+        idx_large = np.argmax(stack_upper_scores[:,1,2]) # Selecting best f measure at 3s
+        score_upper_level = (stack_upper_scores[idx_close,0,:], stack_upper_scores[idx_large,1,:])
+
+        stack_lower_scores = -np.inf * np.ones((len(list_segments_lower_level),2,3))
+        for idx, segments in enumerate(list_segments_lower_level):
+            stack_lower_scores[idx] = self.score_flat_segmentation(segments, dict_annotations, annotation_level = "lower", annotator = annotator)
+        idx_close = np.argmax(stack_lower_scores[:,0,2]) # Selecting best f measure at 0.5s
+        idx_large = np.argmax(stack_lower_scores[:,1,2]) # Selecting best f measure at 3s
+        score_lower_level = (stack_lower_scores[idx_close,0,:], stack_lower_scores[idx_large,1,:])
+
+        return score_upper_level, score_lower_level
+
+
+    def get_sizes_of_annotated_segments(self, annotation_level = "upper", annotator = 1, plot = False):
+        lengths = []
+        for track_id in self.indexes:
+            track = self.all_tracks[track_id]
+
+            try:           
+                # Compute the bars
+                bars = self.get_bars(track.audio_path, index=track_id)
+
+                # Get the annotations
+                dict_annotations = self.get_annotations(track)
+
+                annotations = self.get_this_set_annotations(dict_annotations, annotation_level = annotation_level, annotator = annotator)
+
+                barwise_annot = as_seg.data_manipulation.frontiers_from_time_to_bar(np.array(annotations)[:,1], bars) # Convert the annotations from time to bar
+                for i in range(len(barwise_annot) - 1):
+                    lengths.append(barwise_annot[i+1] - barwise_annot[i]) # Store the length of the annotated segment
+        
+            except FileNotFoundError:
+                print(f'{track_id} not found.')
+                # raise FileNotFoundError(f"Song {track_id} not found, normal ?") from None
+
+        if plot:
+            as_seg.model.current_plot.plot_lenghts_hist(lengths)
+        return lengths
+        
+    # def format_dataset(self, path_audio_files): # TODO
+        # # Copy audio files to the right location.
+        # # Suppose that the audio files are all in the same folder
+        # for track_num in range(len(self.all_tracks)):
+        #     track_idx = self.indexes[track_num]
+        #     song_file_name = self.all_tracks[track_idx].audio_path.split('/')[-1]
+        #     src = f"{path_audio_files}/{song_file_name}" # May change depending on your file structure
+        #     dest = self.all_tracks[track_idx].audio_path
+        #     pathlib.Path(dest).parent.absolute().mkdir(parents=True, exist_ok=True)
+        #     shutil.copy(src, dest)
+    
+
+class BeatlesDataloader(BaseDataloader):
+    def __init__(self, path, feature, cache_path = None, download=False, sr=44100, hop_length = 32, subdivision = 96):
+        super().__init__(feature, cache_path, sr, hop_length, subdivision)
+        self.data_path = path
+        beatles = mirdata.initialize('beatles', data_home = path)
+        if download:
+            beatles.download()            
+        self.all_tracks = beatles.load_tracks()
+        self.indexes = beatles.track_ids
+
+        self.dataset_name = "Beatles"
+
+    def __getitem__(self, index):
+        track_id = self.indexes[index]
+        track = self.all_tracks[track_id]
+
+        # Compute the bars
+        bars = self.get_bars(track.audio_path, index=track_id)
+
+        # Compute the barwise TF matrix
+        barwise_tf_matrix = self.get_barwise_tf_matrix(track.audio_path, bars, index=track_id)
+
+        # Get the annotationks
+        annotations_intervals = track.sections.intervals
+
+        # Return the the bars, the barwise TF matrix and the annotations
+        return track_id, bars, barwise_tf_matrix, annotations_intervals
+    
+    def __len__(self):
+        return len(self.all_tracks)
+    
+    def get_track_of_id(self, track_id):
+        try:
+            index = self.indexes.index(track_id)
+        except ValueError:
+            try:
+                index = self.indexes.index(str(track_id))
+            except ValueError:
+                raise ValueError(f"Track {track_id} not found in the dataset") from None
+        return self.__getitem__(index)
+
+if __name__ == "__main__":
+    # rwcpop = RWCPopDataloader('/home/a23marmo/datasets/rwcpop', feature = "mel", cache_path = "/home/a23marmo/Bureau/data_persisted/rwcpop")
+    # # rwcpop.format_dataset('/home/a23marmo/Bureau/Audio samples/rwcpop/Entire RWC')
+    # for spectrogram, bars, barwise_tf_matrix, track_id, annotations in rwcpop:
+    #     print(spectrogram.shape, track_id)
+
+    salami = SALAMIDataloader('/home/a23marmo/datasets/salami', feature = "mel", cache_path = "/home/a23marmo/Bureau/data_persisted/salami", subset = "train")
+
+    for spectrogram, bars, barwise_tf_matrix, track_id, annotations in salami:
+        try:
+            print(track_id)
+        except FileNotFoundError:
+            print(f'{track_id} not found.')