steerable_retrieval/dataloading/datasets.py

from torch.utils.data import Dataset
from .loading_utils import *
import torch
import os
import random
from tqdm import tqdm
import pandas as pd
from hydra.utils import instantiate

import logging

class BasicProcessor:
    def __init__(self, probability = 1.0, split = None):
        self.probability = probability
        self.split = split
    
    def process(self, annot):
        """
        Process the annotation if the probability is less than the probability and the split is in the split list
        """
        if random.random() < self.probability:
            if self.split is None:
                return self.process(annot)
            if self.split is not None and annot['split'] in self.split:
                return self.process(annot)
            else:
                return annot
        else:
            return annot
        
    def __call__(self, annot):
        return self.process(annot)
    
class RandomNSentencesProcessor(BasicProcessor):

    def process(self, annot):
        
        prompt = annot['prompt']
        sentences = prompt.split('.')
        n_sentences = len(sentences)
        keep_n_sentences = random.randint(1, n_sentences)
        random_sentences = random.sample(random_sentences, keep_n_sentences)
        prompt = '. '.join(random_sentences)
        return prompt
    
class ShuffleSentencesProcessor(BasicProcessor):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
    
    def process(self, annot):
        prompt = annot['prompt']
        sentences = prompt.split('.')
        random.shuffle(sentences)
        prompt = '. '.join(sentences)
        return prompt

class ShuffleTagsProcessor(BasicProcessor):
    def __init__(self, replace_caption_p = 0.5, **kwargs):
        super().__init__(**kwargs)
        self.replace_caption_p = replace_caption_p
    
    def process(self, annot):
        ## if tags is in the annotation, it will be a comma-separated string
        ## drop anywhere from all but 1 to 1 tag, and shuffle
        if 'tags' in annot:
            tags = annot['tags'].split(',')
        else:
            return annot
        len_tags = len(tags)
        drop_tags = random.randint(1, len_tags-1)
        tags_to_drop = random.sample(tags, drop_tags)
        tags = [tag for tag in tags if tag not in tags_to_drop]
        random.shuffle(tags)
        annot['tags'] = '. '.join(tags)

        # if replace caption, replace else add it to the caption
        if random.random() < self.replace_caption_p:
            annot['prompt'] = annot['tags']
        else:
            annot['prompt'] = annot['prompt'] + '.' + annot['tags']
        return annot

class TextAudioDataset(Dataset):
    def __init__(self,
                 annotations = None,
                 get_annotations_function = None,
                 task_kwargs = None,
                 target_n_samples = 96000,
                 target_sr = 48000,
                 return_audio = True,
                 return_text = True,
                 concept = None,
                 return_full_audio = False,
                 preextracted_features = False,
                 truncate_preextracted = 50,
                 split = None,
                 filter_split = None,
                 root_dir = None,
                 new_dir = None,
                 limit_n = None,
                 processors = [],
                 **kwargs
                 ):
        
        
        # Get annotations either directly or from function
        if annotations is not None:
            self.annotations = annotations
        elif get_annotations_function is not None:
            # Support both string and callable
            if isinstance(get_annotations_function, str):
                task_kwargs = task_kwargs or {}
                import importlib

                # parse the fully qualified function name
                module_name, func_name = get_annotations_function.rsplit('.', 1)
                module = importlib.import_module(module_name)
                get_annotations_func = getattr(module, func_name)
                self.annotations = get_annotations_func(**task_kwargs)
            else:
                task_kwargs = task_kwargs or {}
                self.annotations = get_annotations_function(**task_kwargs)
        else:
            raise ValueError("Must provide either annotations or get_annotations_function")
        



        self.target_n_samples = target_n_samples
        self.target_sr = target_sr
        self.return_audio = return_audio
        self.return_text = return_text
        self.concept = concept
        self.return_full_audio = return_full_audio
        self.preextracted_features = preextracted_features
        self.truncate_preextracted = truncate_preextracted
        self.split = split
        self.root_dir = root_dir
        self.new_dir = new_dir
        self.limit_n = limit_n
        # Update split if needed
        if split is not None and split != 'keep':
            for annot in self.annotations:
                annot['split'] = split
        elif split == 'keep':
            # Keep original splits from annotations, or set to 'train' if not present
            for annot in self.annotations:
                if 'split' not in annot or annot['split'] not in ['train', 'val', 'test']:
                    annot['split'] = 'train'
        elif split is None and len(self.annotations) > 0 and 'split' not in self.annotations[0].keys():
            for annot in self.annotations:
                annot['split'] = 'train'

        if filter_split is not None:
            self.annotations = [annot for annot in self.annotations if annot['split'] in filter_split]
                
        annot_df = pd.DataFrame(self.annotations)
        
        try:
            annot_df['file_index'] = pd.factorize(annot_df['file_path'])[0]
        except Exception as e:
            print(e)
        
        annot_df['file_path'] = annot_df['file_path'].apply(lambda x: x.replace(root_dir, new_dir) if root_dir is not None and new_dir is not None else x)
        
        self.annotations = annot_df.to_dict('records')

        # Filter out annotations whose feature/audio file is missing (avoids IndexError when retrying in __getitem__)
        if self.return_audio and self.preextracted_features:
            n_before = len(self.annotations)
            self.annotations = [
                a for a in self.annotations
                if os.path.exists(a['file_path'].replace('.mp3', '.npy').replace('.wav', '.npy'))
            ]
            if len(self.annotations) < n_before:
                logging.warning(
                    f"Filtered out {n_before - len(self.annotations)} annotations with missing feature files. "
                    f"Dataset has {len(self.annotations)} samples."
                )
        elif self.return_audio and not self.preextracted_features:
            n_before = len(self.annotations)
            def _audio_exists(a):
                p = a['file_path']
                return os.path.exists(p.replace('.npy', '.mp3')) or os.path.exists(p.replace('.npy', '.wav'))
            self.annotations = [a for a in self.annotations if _audio_exists(a)]
            if len(self.annotations) < n_before:
                logging.warning(
                    f"Filtered out {n_before - len(self.annotations)} annotations with missing audio files. "
                    f"Dataset has {len(self.annotations)} samples."
                )

        if self.limit_n is not None and self.limit_n < len(self.annotations):
            self.annotations = self.annotations[:self.limit_n]
            print(f"Limiting dataset to {self.limit_n} samples")
        else:
            print(f"Dataset has {len(self.annotations)} samples")

        
        assert return_audio or return_text, "At least one of return_audio or return_text must be True (duh)"

        self.processors = [instantiate(processor) for processor in processors]

        


    def purge(self):
        if self.return_audio and not self.preextracted_features:
            raise NotImplementedError("Purging your audio dataset is probably a bad idea")
        else:
            file_paths = [annot['file_path'] for annot in self.annotations]
            for file_path in file_paths:
                os.remove(file_path)
            print(f"Removed {len(file_paths)} files")
        
    def __len__(self):
        return len(self.annotations)

    def __getitem__(self, idx, return_full_audio = False, hop = None, verbose = False):
        
        
        return_full_audio = self.return_full_audio if return_full_audio is None else return_full_audio
        
        annot = self.annotations[idx]
        
        
        if self.return_audio:
            if not self.preextracted_features:
                # file_path = annot['file_path'].replace('.npy','.mp3').replace('.wav','.mp3')
                # check if the mp3 file exists
                if os.path.exists(annot['file_path'].replace('.npy','.mp3')):
                    file_path = annot['file_path'].replace('.npy','.mp3')
                elif os.path.exists(annot['file_path'].replace('.npy','.wav')):
                    file_path = annot['file_path'].replace('.npy','.wav')
                else:
                    return self.__getitem__(idx+1)
                
                annot['file_path'] = file_path
                try:
                    audio = load_full_and_split(
                        annot['file_path'],
                        self.target_sr,
                        self.target_n_samples,
                        hop=hop,
                        verbose=verbose
                        ) if return_full_audio else load_audio_chunk(
                        annot['file_path'],
                        target_sr=self.target_sr,
                        target_n_samples=self.target_n_samples,
                        verbose=verbose
                        )
                    audio = audio.mean(1)
                except Exception as e:
                    return self.__getitem__(idx+1)
            else:
                file_path = annot['file_path'].replace('.mp3','.npy').replace('.wav','.npy')
                try:

                    audio = np.load(file_path,mmap_mode='r')
                    # Preextracted features may be stored per-clip as a single vector [D]
                    # or as multiple frames [T, D]; sample a random frame when framewise.
                    if audio.ndim > 1:
                        rand_index = random.randint(0, audio.shape[0]-1)
                        audio = audio[rand_index]
                    audio = torch.tensor(np.asarray(audio))
                except Exception as e:
                    return self.__getitem__(idx+1)
                
                
        
        if self.return_text:
            possible_captions = annot['caption']
            # ramdomly choose a caption hash
            
            random_hash = random.choice(list(possible_captions.keys()))
        
            caption = possible_captions[random_hash]
                    
        return_dict = {}
        
        if self.return_audio:
            # Clone so storage is resizable; avoids DataLoader collate error with mmap/numpy-derived tensors
            return_dict['audio'] = audio.clone() if isinstance(audio, torch.Tensor) else audio
            return_dict['file_path'] = annot['file_path']
            
        if self.return_text:
            return_dict['prompt'] = caption
                
        return_dict['file_idx'] = annot['file_index']

        for processor in self.processors:
            return_dict = processor(return_dict)
            
        return return_dict
    
    
    def extract_features(self, model, extract_method = 'extract_features', extract_kwargs = {}, out_key = 'embedding',hop = None, return_full_audio = True, verbose = False):
        
        device = next(model.parameters()).device
        print(f"Extracting features with {extract_method} method on {device} device") if verbose else None
        
        
        for param in model.parameters():
            param.requires_grad = False
        try:
            model.eval()
        except:
            pass
        
        for i in range(len(self)):
            try:
                item = self.__getitem__(i, return_full_audio = return_full_audio, hop = hop, verbose = verbose)
                file_path = self.annotations[i]['file_path'].replace('.mp3','.npy').replace('.wav','.npy')
                
                audio = item['audio'].squeeze(1).to(device)


                
                if audio.shape[0] > 200 :
                    chunks = torch.split(audio, 200, dim=0)
                    chunks = list(chunks)
                    audio_features = []
                    for chunk in chunks:
                        feat = getattr(model, extract_method)(chunk, **extract_kwargs)
                        if out_key is not None:
                            feat = feat[out_key]
                        audio_features.append(feat)
                    audio_features = torch.cat(audio_features, dim=0)
                else:
                    audio_features = getattr(model, extract_method)(audio.to(device), **extract_kwargs)
                    if out_key is not None:
                        audio_features = audio_features[out_key]
                
                    
                print(f"Extracted features for {file_path}, shape: {audio_features.shape}") if verbose else None
                
                yield audio_features, file_path
            except Exception as e:
                print(f"Error extracting features for {file_path}: {e}") if verbose else None
                continue
            
    def extract_and_save_features(self, model, save_dir = None, extract_method = 'extract_features', extract_kwargs = {}, out_key = 'embedding', hop = None, return_full_audio = True, limit_n = None, save = False, verbose = True, root_path = None, done_ids = None):
        
        
        print(self.__len__())
        
        audio_features_all = []
        counter = 0
        skipped_count = 0
        
        save_dir = '' if save_dir is None else save_dir
        done_ids = done_ids or set()
        
        if 's3://' in save_dir:
            import boto3
            import io
            client = boto3.client('s3')
        else:
            client = None
            import io

        # filter self.annotations to only include files that are not in done_ids
        new_annotations = []
        for annot in self.annotations:
            fp = annot['file_path']
            fp = fp.replace(root_path+'/','')
            # remove extension
            fp = fp.replace('.mp3','').replace('.wav','').replace('.npy','')

            if fp not in done_ids:
                new_annotations.append(annot)
        
        self.annotations = new_annotations
        
        for audio_features, file_path in (pbar:= tqdm(self.extract_features(model, extract_method = extract_method, extract_kwargs = extract_kwargs, out_key = out_key, hop = hop, return_full_audio = return_full_audio, verbose = verbose))):
            
            # print(file_path, root_path, save_dir)
            
            if root_path is not None:
                file_path = file_path.replace(root_path+'/','')

            save_path = os.path.join(save_dir, file_path)
            
            if save and audio_features is not None:
                

                #remove the root path from the file path
                
                
                if 's3://' in save_dir:
                    bucket, key = save_dir.replace("s3://", "").split("/", 1)
                    key = f"{key}/{file_path}"
                    
                    # local_path = os.path.join(local_temp_dir, file_path)
                    # os.makedirs(os.path.dirname(local_path), exist_ok=True)
                    # np.save(local_path, audio_features.detach().cpu().numpy())
                    
                    pbar.set_description(f"Uploading features to s3://{bucket}/{key}") if verbose else None
                    try:
                        # client.upload_file(save_path, bucket, key)
                        
                        buffer = io.BytesIO()
                        np.save(buffer, audio_features.detach().cpu().numpy())
                        buffer.seek(0)
                        client.put_object(Bucket=bucket, Key=key, Body=buffer)
                    except Exception as e:
                        print(f"Error uploading to s3: {e}") if verbose else None
                        
                    # os.remove(local_path)
                else:
                    pbar.set_description(f"Saving features in {save_path}, shape: {audio_features.shape}")
                    os.makedirs(os.path.dirname(save_path), exist_ok=True)
                    if os.path.exists(save_path):
                        os.remove(save_path)
                    np.save(save_path, audio_features.detach().cpu().numpy())
            
            if not save and audio_features is not None:
                pbar.set_description(f"{file_path}, shape: {audio_features.shape}")
                pass
                
            audio_features_all.append(audio_features.detach().cpu()) if audio_features is not None else None
            
            counter += 1
            if limit_n and counter >= limit_n:
                break
        
        if skipped_count > 0:
            print(f"Skipped {skipped_count} already processed items") if verbose else None
        
        try:   
            print(f"Returning {len(audio_features_all)} features") if verbose else None
            all_= torch.stack(audio_features_all)
            print(f"Stacked features, shape: {all_.shape}") if verbose else None
            return all_
        
        
        except:
            return None