import contextlib import io import json import logging import os.path import random import re import time import pandas as pd from a_cls.dataloader import make_midname_dict from open_clip import get_tokenizer from open_clip.factory import HF_HUB_PREFIX from .process_video import load_and_transform_video, get_video_transform from .process_audio import load_and_transform_audio, get_audio_transform from .process_text import load_and_transform_text from .process_depth import load_and_transform_depth, get_depth_transform from .process_thermal import load_and_transform_thermal, get_thermal_transform import argparse from os.path import join as opj from torch.utils.data import Dataset, DataLoader from tqdm import tqdm class VAT_dataset(Dataset): def __init__(self, args): super().__init__() self.video_decode_backend = args.video_decode_backend self.num_frames = args.num_frames self.text_type = args.text_type self.total_text = ['raw', 'mplug', 'polish_mplug', 'sound_mplug'] + [f'ofa{i}' for i in range(8)] self.weight = [0.2, 0.2, 0.2, 0.2] + [0.2 / 8] * 8 self.title = self.text_type == 'raw' self.data_root = '/apdcephfs_cq3/share_1311970/A_Youtube' if args.clip_type != 'al': with open(args.train_data, 'r') as f: self.id2title_folder_caps = json.load(f) self.ids = list(self.id2title_folder_caps.keys())[:args.train_num_samples] else: self.id2path_cap, self.ids = get_audio_anno() self.clip_type = args.clip_type self.num_mel_bins = args.num_mel_bins self.target_length = args.target_length self.audio_sample_rate = args.audio_sample_rate self.audio_mean = args.audio_mean self.audio_std = args.audio_std # self.audio_error_file = open('./audio_error_id.txt', 'w') self.tokenizer = get_tokenizer(HF_HUB_PREFIX + args.model, cache_dir=args.cache_dir) self.video_transform = get_video_transform(args) self.audio_transform = get_audio_transform(args) self.depth_transform = get_depth_transform(args) self.thermal_transform = get_thermal_transform(args) def __len__(self): return len(self.ids) # return self.id2title_folder_caps.shape[0] def __getitem__(self, idx): try: if self.clip_type == 'al': matched_modality, input_ids, attention_mask = self.get_audio_text(idx) return matched_modality, input_ids, attention_mask else: id = self.ids[idx] folder = self.id2title_folder_caps[id]['folder'] text_output, ofa_number = self.get_text(id) input_ids, attention_mask = text_output['input_ids'], text_output['attention_mask'] if self.clip_type == 'vl' or self.clip_type == 'vl_new': matched_modality = self.get_video(id, folder) # elif self.clip_type == 'al': # matched_modality = self.get_audio(id, folder) elif self.clip_type == 'dl': matched_modality = self.get_depth(id, folder, ofa_number) elif self.clip_type == 'tl': matched_modality = self.get_thermal(id, folder, ofa_number) return matched_modality['pixel_values'], input_ids, attention_mask except Exception as error_msg: logging.info(f"Failed at {idx} with \"{error_msg}\"") return self.__getitem__(random.randint(0, self.__len__()-1)) def get_video(self, id, folder): # video_path = opj(self.data_root, folder, f'{id}.mp4') resize_folder = 'new_download_resize256_skip15' if folder.startswith('new_') else f'{folder}_resize256_skip15' video_path = opj(self.data_root, resize_folder, f'{id}.mp4') video = load_and_transform_video(video_path, self.video_transform, video_decode_backend=self.video_decode_backend, num_frames=self.num_frames) return video def get_audio_text(self, idx): path_cap = self.id2path_cap[self.ids[idx]] audio_path = path_cap['path'] audio_data = load_and_transform_audio(audio_path, self.audio_transform) caption = path_cap['caption'] if isinstance(caption, list): if isinstance(caption[0], str) and len(caption) > 1: caption = random.choice(caption) else: caption = caption[0] input_ids, attention_mask = self.tokenizer(caption) return audio_data, input_ids.squeeze(), attention_mask.squeeze() # def get_audio(self, idx): ''' audio_path = opj(self.data_root, folder, f'{id}.mp3') if os.path.exists(audio_path): pass else: audio_path = audio_path[:-4] + '.m4a' if os.path.exists(audio_path): pass else: audio_path = audio_path[:-4] + '.wav' if not os.path.exists(audio_path): # self.audio_error_file.write(audio_path[:-4] + '\n') raise FileNotFoundError(f'Not found audio file at \'{audio_path[:-4]}\' with .mp3 .m4a .wav') # AudioSegment.from_file(audio_path).export(audio_path[:-4] + '.mp3', format='mp3') # audio_path = opj(self.data_root, folder, f'{id}.mp3') audio = load_and_transform_audio(audio_path, self.audio_transform) ''' # audio_path = opj(self.data_root, folder+'_ffmpeg_mp3', f'{id}.mp3') # audio = load_and_transform_audio(audio_path, self.audio_transform) ''' audiocap_id = self.meta['uniq_id'][idx] audio_path = f'/apdcephfs_cq3/share_1311970/downstream_datasets/Audio/audiocaps/audio/train/{audiocap_id}.flac' audio_data = load_and_transform_audio(audio_path, self.audio_transform) caption = self.meta['text'][idx] input_ids, attention_mask = self.tokenizer(caption) return audio_data, input_ids.squeeze(), attention_mask.squeeze() ''' ''' path_cap = self.id2path_cap[self.ids[idx]] audio_path = f"/remote-home/freesound/{path_cap['path']}" audio_data = load_and_transform_audio(audio_path, self.audio_transform) caption = path_cap['caption'] input_ids, attention_mask = self.tokenizer(caption) ''' # return audio def get_text(self, id): if self.text_type != 'mix': text = self.id2title_folder_caps[id][self.text_type] text_output = load_and_transform_text(text, self.tokenizer, title=self.title) return text_output, None else: text_type = random.choices(self.total_text, self.weight)[0] ofa_number = None if text_type.startswith('ofa'): ofa_number = int(text_type[-1]) text = self.id2title_folder_caps[id]['ofa'][ofa_number] else: text = self.id2title_folder_caps[id][text_type] text_output = load_and_transform_text(text, self.tokenizer, title=text_type=='raw') return text_output, ofa_number def get_depth(self, id, folder, ofa_number): depth_folder = opj(self.data_root, folder, f'{id}_depth_f8glpn_folder') random_id = random.randint(0, 7) if ofa_number is None else ofa_number # random_id = 3 depth_path = os.path.join(depth_folder, f'{random_id}.png') depth = load_and_transform_depth(depth_path, self.depth_transform) return depth def get_thermal(self, id, folder, ofa_number): thermal_folder = opj(self.data_root, folder, f'{id}_thermal_folder') random_id = random.randint(0, 7) if ofa_number is None else ofa_number # random_id = 3 thermal_path = os.path.join(thermal_folder, f'{random_id}.jpg') thermal = load_and_transform_thermal(thermal_path, self.thermal_transform) return thermal if __name__ == '__main__': parser = argparse.ArgumentParser('Pre-training', add_help=False) parser.add_argument('--num_frames', default=8, type=float, help='') parser.add_argument('--workers', default=10, type=int, help='') args = parser.parse_args() args.cache_dir = 'D:\Omni-modal-hf' args.num_frames = 8 args.clip_type = 'vl' args.num_mel_bins = 128 args.target_length = 1024 args.audio_sample_rate = 16000 args.audio_mean = 1 args.audio_std = 1 args.rank = 0 args.batch_size = 16 train_dataset = VAT_dataset(args) load = DataLoader(train_dataset, batch_size=args.batch_size, num_workers=args.workers) for samples in tqdm((load)): matched_modality, input_ids, attention_mask = samples # print(video.shape, text.shape)