| """Ref-AVS dataset: frames, masks, log-mel audio, and referring expressions.""" |
| import os |
| import numpy |
| import torch |
| import pandas |
|
|
| from dataloader.visual.visual_dataset import Visual |
| from dataloader.audio.audio_and_text_dataset import AudioAndText |
|
|
|
|
| class AV(torch.utils.data.Dataset): |
| """Pairs ``Visual`` with ``AudioAndText`` via REFAVS ``metadata.csv``.""" |
|
|
| def __init__(self, split, augmentation, param, root_path=''): |
| self.visual_dataset = Visual( |
| augmentation['visual'], root_path, split, |
| param.image_size, param.image_embedding_size, |
| ) |
| self.audio_and_text_dataset = AudioAndText(augmentation['audio'], root_path, split) |
| self.split = split |
| self.file_path = self.organise_files(self.split, root_path, csv_name_='metadata.csv') |
|
|
| def __getitem__(self, index): |
| vid, fid, exp, _ = self.file_path[index] |
| frame, label, prompts = self.visual_dataset.load_data(vid, fid) |
| audio_mel, text_feature = self.audio_and_text_dataset.load_audio_wave(vid, exp) |
| return { |
| 'frame': frame, |
| 'label': label, |
| 'spectrogram': audio_mel, |
| 'text': text_feature, |
| 'id': self.file_path[index], |
| 'prompts': prompts, |
| } |
|
|
| def __len__(self): |
| return len(self.file_path) |
|
|
| @staticmethod |
| def organise_files(split_, root_path_, csv_name_): |
| total_files = pandas.read_csv(os.path.join(root_path_, csv_name_)) |
| if split_ == 'test_n': |
| rows = zip( |
| total_files[total_files['split'] == split_]['uid'], |
| total_files[total_files['split'] == split_]['fid'], |
| total_files[total_files['split'] == split_]['exp'], |
| ) |
| return [ |
| [name.rsplit('_', 2)[0], object_id, expression, 0] |
| for name, object_id, expression in rows |
| ] |
|
|
| rows = zip( |
| total_files[total_files['split'] == split_]['vid'], |
| total_files[total_files['split'] == split_]['fid'], |
| total_files[total_files['split'] == split_]['exp'], |
| ) |
| file_path = [[vid, fid, expression, 0] for vid, fid, expression in rows] |
|
|
| if split_ == 'train': |
| null_uids = list(total_files[total_files['split'] == split_]['uid']) |
| assert len(null_uids) == len(file_path) |
| for idx, row in enumerate(file_path): |
| if 'null_' in null_uids[idx]: |
| row[0] = null_uids[idx].rsplit('_', 2)[0] |
| row[-1] = null_uids[idx].rsplit('_', 2)[1] |
|
|
| return file_path |
|
|
