ldm/data/tsvdataset.py

from glob import glob
from torch.utils.data import Dataset
import numpy as np
import pandas as pd

class TSVDataset(Dataset):
 def __init__(self, tsv_path, spec_crop_len=None):
 super().__init__()
 self.batch_max_length = spec_crop_len
 self.batch_min_length = 50
 df = pd.read_csv(tsv_path,sep='\t')
 df = self.add_name_num(df)
 self.dataset = df
 print('dataset len:', len(self.dataset))

 def add_name_num(self,df):
 """each file may have different caption, we add num to filename to identify each audio-caption pair"""
 name_count_dict = {}
 change = []
 for t in df.itertuples():
 name = getattr(t,'name')
 if name in name_count_dict:
 name_count_dict[name] += 1
 else:
 name_count_dict[name] = 0
 change.append((t[0],name_count_dict[name]))
 for t in change:
 df.loc[t[0],'name'] = df.loc[t[0],'name'] + f'_{t[1]}'
 return df


 def __getitem__(self, idx):
 data = self.dataset.iloc[idx]
 item = {}
 spec = np.load(data['mel_path']) # mel spec [80, 624]
 if spec.shape[1] <= self.batch_max_length:
 spec = np.pad(spec, ((0, 0), (0, self.batch_max_length - spec.shape[1]))) # [80, 624]

 item['image'] = spec
 item["caption"] = data['caption']
 item["f_name"] = data['name']
 return item

 def __len__(self):
 return len(self.dataset)

class TSVDatasetStruct(TSVDataset):
 def __getitem__(self, idx):
 data = self.dataset.iloc[idx]
 item = {}
 spec = np.load(data['mel_path']) # mel spec [80, 624]
 if spec.shape[1] <= self.batch_max_length:
 spec = np.pad(spec, ((0, 0), (0, self.batch_max_length - spec.shape[1]))) # [80, 624]

 item['image'] = spec[:,:self.batch_max_length]
 item["caption"] = {'ori_caption':data['ori_cap'],'struct_caption':data['caption']}
 item["f_name"] = data['name']
 return item

class TSVDatasetTestFake(TSVDataset):
 def __init__(self, specs_dataset_cfg):
 super().__init__(phase='test', **specs_dataset_cfg)
 self.dataset = [self.dataset[0]]