import csv
import os
import pickle
import sys
import numpy as np
import torch
import random
import math
import librosa
class audio_video_spec_fullset_Dataset(torch.utils.data.Dataset):
# Only Load audio dataset: for training Stage1: Audio Npy Dataset
def __init__(self, split, dataset1, feat_type='clip', transforms=None, sr=22050, duration=10, truncate=220000, fps=21.5, drop=0.0, fix_frames=False, hop_len=256):
super().__init__()
if split == "train":
self.split = "Train"
elif split == "valid" or split == 'test':
self.split = "Test"
# Default params:
self.min_duration = 2
self.sr = sr # 22050
self.duration = duration # 10
self.truncate = truncate # 220000
self.fps = fps
self.fix_frames = fix_frames
self.hop_len = hop_len
self.drop = drop
print("Fix Frames: {}".format(self.fix_frames))
print("Use Drop: {}".format(self.drop))
# Dataset1: (VGGSound)
assert dataset1.dataset_name == "VGGSound"
# spec_dir: spectrogram path
# feat_dir: CAVP feature path
# video_dir: video path
dataset1_spec_dir = os.path.join(dataset1.data_dir, "mel_maa2", "npy")
dataset1_feat_dir = os.path.join(dataset1.data_dir, "cavp")
dataset1_video_dir = os.path.join(dataset1.video_dir, "tmp_vid")
split_txt_path = dataset1.split_txt_path
with open(os.path.join(split_txt_path, '{}.txt'.format(self.split)), "r") as f:
data_list1 = f.readlines()
data_list1 = list(map(lambda x: x.strip(), data_list1))
spec_list1 = list(map(lambda x: os.path.join(dataset1_spec_dir, x) + "_mel.npy", data_list1)) # spec
feat_list1 = list(map(lambda x: os.path.join(dataset1_feat_dir, x) + ".npz", data_list1)) # feat
video_list1 = list(map(lambda x: os.path.join(dataset1_video_dir, x) + "_new_fps_21.5_truncate_0_10.0.mp4", data_list1)) # video
# Merge Data:
self.data_list = data_list1 if self.split != "Test" else data_list1[:200]
self.spec_list = spec_list1 if self.split != "Test" else spec_list1[:200]
self.feat_list = feat_list1 if self.split != "Test" else feat_list1[:200]
self.video_list = video_list1 if self.split != "Test" else video_list1[:200]
assert len(self.data_list) == len(self.spec_list) == len(self.feat_list) == len(self.video_list)
shuffle_idx = np.random.permutation(np.arange(len(self.data_list)))
self.data_list = [self.data_list[i] for i in shuffle_idx]
self.spec_list = [self.spec_list[i] for i in shuffle_idx]
self.feat_list = [self.feat_list[i] for i in shuffle_idx]
self.video_list = [self.video_list[i] for i in shuffle_idx]
print('Split: {} Sample Num: {}'.format(split, len(self.data_list)))
def __len__(self):
return len(self.data_list)
def load_spec_and_feat(self, spec_path, video_feat_path):
"""Load audio spec and video feat"""
try:
spec_raw = np.load(spec_path).astype(np.float32) # channel: 1
except:
print(f"corrupted mel: {spec_path}", flush=True)
spec_raw = np.zeros((80, 625), dtype=np.float32) # [C, T]
p = np.random.uniform(0,1)
if p > self.drop:
try:
video_feat = np.load(video_feat_path)['feat'].astype(np.float32)
except:
print(f"corrupted video: {video_feat_path}", flush=True)
video_feat = np.load(os.path.join(os.path.dirname(video_feat_path), 'empty_vid.npz'))['feat'].astype(np.float32)
else:
video_feat = np.load(os.path.join(os.path.dirname(video_feat_path), 'empty_vid.npz'))['feat'].astype(np.float32)
spec_len = self.sr * self.duration / self.hop_len
if spec_raw.shape[1] < spec_len:
spec_raw = np.tile(spec_raw, math.ceil(spec_len / spec_raw.shape[1]))
spec_raw = spec_raw[:, :int(spec_len)]
feat_len = self.fps * self.duration
if video_feat.shape[0] < feat_len:
video_feat = np.tile(video_feat, (math.ceil(feat_len / video_feat.shape[0]), 1))
video_feat = video_feat[:int(feat_len)]
return spec_raw, video_feat
def mix_audio_and_feat(self, spec1=None, spec2=None, video_feat1=None, video_feat2=None, video_info_dict={}, mode='single'):
""" Return Mix Spec and Mix video feat"""
if mode == "single":
# spec1:
if not self.fix_frames:
start_idx = random.randint(0, self.sr * self.duration - self.truncate - 1) # audio start
else:
start_idx = 0
start_frame = int(self.fps * start_idx / self.sr)
truncate_frame = int(self.fps * self.truncate / self.sr)
# Spec Start & Truncate:
spec_start = int(start_idx / self.hop_len)
spec_truncate = int(self.truncate / self.hop_len)
spec1 = spec1[:, spec_start : spec_start + spec_truncate]
video_feat1 = video_feat1[start_frame: start_frame + truncate_frame]
# info_dict:
video_info_dict['video_time1'] = str(start_frame) + '_' + str(start_frame+truncate_frame) # Start frame, end frame
video_info_dict['video_time2'] = ""
return spec1, video_feat1, video_info_dict
elif mode == "concat":
total_spec_len = int(self.truncate / self.hop_len)
# Random Trucate len:
spec1_truncate_len = random.randint(self.min_duration * self.sr // self.hop_len, total_spec_len - self.min_duration * self.sr // self.hop_len - 1)
spec2_truncate_len = total_spec_len - spec1_truncate_len
# Sample spec clip:
spec_start1 = random.randint(0, total_spec_len - spec1_truncate_len - 1)
spec_start2 = random.randint(0, total_spec_len - spec2_truncate_len - 1)
spec_end1, spec_end2 = spec_start1 + spec1_truncate_len, spec_start2 + spec2_truncate_len
# concat spec:
spec1, spec2 = spec1[:, spec_start1 : spec_end1], spec2[:, spec_start2 : spec_end2]
concat_audio_spec = np.concatenate([spec1, spec2], axis=1)
# Concat Video Feat:
start1_frame, truncate1_frame = int(self.fps * spec_start1 * self.hop_len / self.sr), int(self.fps * spec1_truncate_len * self.hop_len / self.sr)
start2_frame, truncate2_frame = int(self.fps * spec_start2 * self.hop_len / self.sr), int(self.fps * self.truncate / self.sr) - truncate1_frame
video_feat1, video_feat2 = video_feat1[start1_frame : start1_frame + truncate1_frame], video_feat2[start2_frame : start2_frame + truncate2_frame]
concat_video_feat = np.concatenate([video_feat1, video_feat2])
video_info_dict['video_time1'] = str(start1_frame) + '_' + str(start1_frame+truncate1_frame) # Start frame, end frame
video_info_dict['video_time2'] = str(start2_frame) + '_' + str(start2_frame+truncate2_frame)
return concat_audio_spec, concat_video_feat, video_info_dict
def __getitem__(self, idx):
audio_name1 = self.data_list[idx]
spec_npy_path1 = self.spec_list[idx]
video_feat_path1 = self.feat_list[idx]
video_path1 = self.video_list[idx]
# select other video:
flag = False
if random.uniform(0, 1) < 0.5:
flag = True
random_idx = idx
while random_idx == idx:
random_idx = random.randint(0, len(self.data_list)-1)
audio_name2 = self.data_list[random_idx]
spec_npy_path2 = self.spec_list[random_idx]
video_feat_path2 = self.feat_list[random_idx]
video_path2 = self.video_list[random_idx]
# Load the Spec and Feat:
spec1, video_feat1 = self.load_spec_and_feat(spec_npy_path1, video_feat_path1)
if flag:
spec2, video_feat2 = self.load_spec_and_feat(spec_npy_path2, video_feat_path2)
video_info_dict = {'audio_name1':audio_name1, 'audio_name2': audio_name2, 'video_path1': video_path1, 'video_path2': video_path2}
mix_spec, mix_video_feat, mix_info = self.mix_audio_and_feat(spec1, spec2, video_feat1, video_feat2, video_info_dict, mode='concat')
else:
video_info_dict = {'audio_name1':audio_name1, 'audio_name2': "", 'video_path1': video_path1, 'video_path2': ""}
mix_spec, mix_video_feat, mix_info = self.mix_audio_and_feat(spec1=spec1, video_feat1=video_feat1, video_info_dict=video_info_dict, mode='single')
# print("mix spec shape:", mix_spec.shape)
# print("mix video feat:", mix_video_feat.shape)
data_dict = {}
# data_dict['mix_spec'] = mix_spec[None].repeat(3, axis=0) # TODO：要把这里改掉，否则无法适应maa的autoencoder
data_dict['mix_spec'] = mix_spec # (80, 512)
data_dict['mix_video_feat'] = mix_video_feat # (32, 512)
data_dict['mix_info_dict'] = mix_info
return data_dict
class audio_video_spec_fullset_Dataset_Train(audio_video_spec_fullset_Dataset):
def __init__(self, dataset_cfg):
super().__init__(split='train', **dataset_cfg)
class audio_video_spec_fullset_Dataset_Valid(audio_video_spec_fullset_Dataset):
def __init__(self, dataset_cfg):
super().__init__(split='valid', **dataset_cfg)
class audio_video_spec_fullset_Dataset_Test(audio_video_spec_fullset_Dataset):
def __init__(self, dataset_cfg):
super().__init__(split='test', **dataset_cfg)
class audio_video_spec_fullset_Dataset_inpaint(audio_video_spec_fullset_Dataset):
def __getitem__(self, idx):
audio_name1 = self.data_list[idx]
spec_npy_path1 = self.spec_list[idx]
video_feat_path1 = self.feat_list[idx]
video_path1 = self.video_list[idx]
# Load the Spec and Feat:
spec1, video_feat1 = self.load_spec_and_feat(spec_npy_path1, video_feat_path1)
video_info_dict = {'audio_name1': audio_name1, 'audio_name2': "", 'video_path1': video_path1, 'video_path2': ""}
mix_spec, mix_masked_spec, mix_video_feat, mix_info = self.mix_audio_and_feat(spec1=spec1, video_feat1=video_feat1, video_info_dict=video_info_dict)
# print("mix spec shape:", mix_spec.shape)
# print("mix video feat:", mix_video_feat.shape)
data_dict = {}
# data_dict['mix_spec'] = mix_spec[None].repeat(3, axis=0) # TODO：要把这里改掉，否则无法适应maa的autoencoder
data_dict['mix_spec'] = mix_spec # (80, 512)
data_dict['hybrid_feat'] = {'mix_video_feat': mix_video_feat, 'mix_spec': mix_masked_spec} # (32, 512)
data_dict['mix_info_dict'] = mix_info
return data_dict
def mix_audio_and_feat(self, spec1=None, video_feat1=None, video_info_dict={}):
""" Return Mix Spec and Mix video feat"""
# spec1:
if not self.fix_frames:
start_idx = random.randint(0, self.sr * self.duration - self.truncate - 1) # audio start
else:
start_idx = 0
start_frame = int(self.fps * start_idx / self.sr)
truncate_frame = int(self.fps * self.truncate / self.sr)
# Spec Start & Truncate:
spec_start = int(start_idx / self.hop_len)
spec_truncate = int(self.truncate / self.hop_len)
spec1 = spec1[:, spec_start: spec_start + spec_truncate]
video_feat1 = video_feat1[start_frame: start_frame + truncate_frame]
# Start masking frames:
masked_spec = random.randint(1, int(spec_truncate * 0.5 // 16)) * 16 # 16帧的倍数，最多mask 50%
masked_truncate = int(masked_spec * self.hop_len)
masked_frame = int(self.fps * masked_truncate / self.sr)
start_masked_idx = random.randint(0, self.truncate - masked_truncate - 1)
start_masked_frame = int(self.fps * start_masked_idx / self.sr)
start_masked_spec = int(start_masked_idx / self.hop_len)
masked_spec1 = np.zeros((80, spec_truncate)).astype(np.float32)
masked_spec1[:] = spec1[:]
masked_spec1[:, start_masked_spec:start_masked_spec+masked_spec] = np.zeros((80, masked_spec))
video_feat1[start_masked_frame:start_masked_frame+masked_frame, :] = np.zeros((masked_frame, 512))
# info_dict:
video_info_dict['video_time1'] = str(start_frame) + '_' + str(start_frame + truncate_frame) # Start frame, end frame
video_info_dict['video_time2'] = ""
return spec1, masked_spec1, video_feat1, video_info_dict
class audio_video_spec_fullset_Dataset_inpaint_Train(audio_video_spec_fullset_Dataset_inpaint):
def __init__(self, dataset_cfg):
super().__init__(split='train', **dataset_cfg)
class audio_video_spec_fullset_Dataset_inpaint_Valid(audio_video_spec_fullset_Dataset_inpaint):
def __init__(self, dataset_cfg):
super().__init__(split='valid', **dataset_cfg)
class audio_video_spec_fullset_Dataset_inpaint_Test(audio_video_spec_fullset_Dataset_inpaint):
def __init__(self, dataset_cfg):
super().__init__(split='test', **dataset_cfg)
class audio_Dataset(torch.utils.data.Dataset):
# Only Load audio dataset: for training Stage1: Audio Npy Dataset
def __init__(self, split, dataset1, sr=22050, duration=10, truncate=220000, debug_num=False, fix_frames=False, hop_len=256):
super().__init__()
if split == "train":
self.split = "Train"
elif split == "valid" or split == 'test':
self.split = "Test"
# Default params:
self.min_duration = 2
self.sr = sr # 22050
self.duration = duration # 10
self.truncate = truncate # 220000
self.fix_frames = fix_frames
self.hop_len = hop_len
print("Fix Frames: {}".format(self.fix_frames))
# Dataset1: (VGGSound)
assert dataset1.dataset_name == "VGGSound"
# spec_dir: spectrogram path
# dataset1_spec_dir = os.path.join(dataset1.data_dir, "codec")
dataset1_wav_dir = os.path.join(dataset1.wav_dir, "wav")
split_txt_path = dataset1.split_txt_path
with open(os.path.join(split_txt_path, '{}.txt'.format(self.split)), "r") as f:
data_list1 = f.readlines()
data_list1 = list(map(lambda x: x.strip(), data_list1))
wav_list1 = list(map(lambda x: os.path.join(dataset1_wav_dir, x) + ".wav", data_list1)) # feat
# Merge Data:
self.data_list = data_list1
self.wav_list = wav_list1
assert len(self.data_list) == len(self.wav_list)
shuffle_idx = np.random.permutation(np.arange(len(self.data_list)))
self.data_list = [self.data_list[i] for i in shuffle_idx]
self.wav_list = [self.wav_list[i] for i in shuffle_idx]
if debug_num:
self.data_list = self.data_list[:debug_num]
self.wav_list = self.wav_list[:debug_num]
print('Split: {} Sample Num: {}'.format(split, len(self.data_list)))
def __len__(self):
return len(self.data_list)
def load_spec_and_feat(self, wav_path):
"""Load audio spec and video feat"""
try:
wav_raw, sr = librosa.load(wav_path, sr=self.sr) # channel: 1
except:
print(f"corrupted wav: {wav_path}", flush=True)
wav_raw = np.zeros((160000,), dtype=np.float32) # [T]
wav_len = self.sr * self.duration
if wav_raw.shape[0] < wav_len:
wav_raw = np.tile(wav_raw, math.ceil(wav_len / wav_raw.shape[0]))
wav_raw = wav_raw[:int(wav_len)]
return wav_raw
def mix_audio_and_feat(self, wav_raw1=None, video_info_dict={}, mode='single'):
""" Return Mix Spec and Mix video feat"""
if mode == "single":
# spec1:
if not self.fix_frames:
start_idx = random.randint(0, self.sr * self.duration - self.truncate - 1) # audio start
else:
start_idx = 0
wav_start = start_idx
wav_truncate = self.truncate
wav_raw1 = wav_raw1[wav_start: wav_start + wav_truncate]
return wav_raw1, video_info_dict
elif mode == "concat":
total_spec_len = int(self.truncate / self.hop_len)
# Random Trucate len:
spec1_truncate_len = random.randint(self.min_duration * self.sr // self.hop_len, total_spec_len - self.min_duration * self.sr // self.hop_len - 1)
spec2_truncate_len = total_spec_len - spec1_truncate_len
# Sample spec clip:
spec_start1 = random.randint(0, total_spec_len - spec1_truncate_len - 1)
spec_start2 = random.randint(0, total_spec_len - spec2_truncate_len - 1)
spec_end1, spec_end2 = spec_start1 + spec1_truncate_len, spec_start2 + spec2_truncate_len
# concat spec:
return video_info_dict
def __getitem__(self, idx):
audio_name1 = self.data_list[idx]
wav_path1 = self.wav_list[idx]
# select other video:
flag = False
if random.uniform(0, 1) < -1:
flag = True
random_idx = idx
while random_idx == idx:
random_idx = random.randint(0, len(self.data_list)-1)
audio_name2 = self.data_list[random_idx]
spec_npy_path2 = self.spec_list[random_idx]
wav_path2 = self.wav_list[random_idx]
# Load the Spec and Feat:
wav_raw1 = self.load_spec_and_feat(wav_path1)
if flag:
spec2, video_feat2 = self.load_spec_and_feat(spec_npy_path2, wav_path2)
video_info_dict = {'audio_name1':audio_name1, 'audio_name2': audio_name2}
mix_spec, mix_video_feat, mix_info = self.mix_audio_and_feat(video_info_dict, mode='concat')
else:
video_info_dict = {'audio_name1':audio_name1, 'audio_name2': ""}
mix_wav, mix_info = self.mix_audio_and_feat(wav_raw1=wav_raw1, video_info_dict=video_info_dict, mode='single')
data_dict = {}
data_dict['mix_wav'] = mix_wav # (131072,)
data_dict['mix_info_dict'] = mix_info
return data_dict
class audio_Dataset_Train(audio_Dataset):
def __init__(self, dataset_cfg):
super().__init__(split='train', **dataset_cfg)
class audio_Dataset_Test(audio_Dataset):
def __init__(self, dataset_cfg):
super().__init__(split='test', **dataset_cfg)
class audio_Dataset_Valid(audio_Dataset):
def __init__(self, dataset_cfg):
super().__init__(split='valid', **dataset_cfg)
class video_codec_Dataset(torch.utils.data.Dataset):
# Only Load audio dataset: for training Stage1: Audio Npy Dataset
def __init__(self, split, dataset1, sr=22050, duration=10, truncate=220000, fps=21.5, debug_num=False, fix_frames=False, hop_len=256):
super().__init__()
if split == "train":
self.split = "Train"
elif split == "valid" or split == 'test':
self.split = "Test"
# Default params:
self.min_duration = 2
self.fps = fps
self.sr = sr # 22050
self.duration = duration # 10
self.truncate = truncate # 220000
self.fix_frames = fix_frames
self.hop_len = hop_len
print("Fix Frames: {}".format(self.fix_frames))
# Dataset1: (VGGSound)
assert dataset1.dataset_name == "VGGSound"
# spec_dir: spectrogram path
# dataset1_spec_dir = os.path.join(dataset1.data_dir, "codec")
dataset1_feat_dir = os.path.join(dataset1.data_dir, "cavp")
dataset1_wav_dir = os.path.join(dataset1.wav_dir, "wav")
split_txt_path = dataset1.split_txt_path
with open(os.path.join(split_txt_path, '{}.txt'.format(self.split)), "r") as f:
data_list1 = f.readlines()
data_list1 = list(map(lambda x: x.strip(), data_list1))
wav_list1 = list(map(lambda x: os.path.join(dataset1_wav_dir, x) + ".wav", data_list1)) # feat
feat_list1 = list(map(lambda x: os.path.join(dataset1_feat_dir, x) + ".npz", data_list1)) # feat
# Merge Data:
self.data_list = data_list1
self.wav_list = wav_list1
self.feat_list = feat_list1
assert len(self.data_list) == len(self.wav_list)
shuffle_idx = np.random.permutation(np.arange(len(self.data_list)))
self.data_list = [self.data_list[i] for i in shuffle_idx]
self.wav_list = [self.wav_list[i] for i in shuffle_idx]
self.feat_list = [self.feat_list[i] for i in shuffle_idx]
if debug_num:
self.data_list = self.data_list[:debug_num]
self.wav_list = self.wav_list[:debug_num]
self.feat_list = self.feat_list[:debug_num]
print('Split: {} Sample Num: {}'.format(split, len(self.data_list)))
def __len__(self):
return len(self.data_list)
def load_spec_and_feat(self, wav_path, video_feat_path):
"""Load audio spec and video feat"""
try:
wav_raw, sr = librosa.load(wav_path, sr=self.sr) # channel: 1
except:
print(f"corrupted wav: {wav_path}", flush=True)
wav_raw = np.zeros((160000,), dtype=np.float32) # [T]
try:
video_feat = np.load(video_feat_path)['feat'].astype(np.float32)
except:
print(f"corrupted video: {video_feat_path}", flush=True)
video_feat = np.load(os.path.join(os.path.dirname(video_feat_path), 'empty_vid.npz'))['feat'].astype(np.float32)
wav_len = self.sr * self.duration
if wav_raw.shape[0] < wav_len:
wav_raw = np.tile(wav_raw, math.ceil(wav_len / wav_raw.shape[0]))
wav_raw = wav_raw[:int(wav_len)]
feat_len = self.fps * self.duration
if video_feat.shape[0] < feat_len:
video_feat = np.tile(video_feat, (math.ceil(feat_len / video_feat.shape[0]), 1))
video_feat = video_feat[:int(feat_len)]
return wav_raw, video_feat
def mix_audio_and_feat(self, wav_raw1=None, video_feat1=None, video_info_dict={}, mode='single'):
""" Return Mix Spec and Mix video feat"""
if mode == "single":
# spec1:
if not self.fix_frames:
start_idx = random.randint(0, self.sr * self.duration - self.truncate - 1) # audio start
else:
start_idx = 0
wav_start = start_idx
wav_truncate = self.truncate
wav_raw1 = wav_raw1[wav_start: wav_start + wav_truncate]
start_frame = int(self.fps * start_idx / self.sr)
truncate_frame = int(self.fps * self.truncate / self.sr)
video_feat1 = video_feat1[start_frame: start_frame + truncate_frame]
# info_dict:
video_info_dict['video_time1'] = str(start_frame) + '_' + str(start_frame+truncate_frame) # Start frame, end frame
video_info_dict['video_time2'] = ""
return wav_raw1, video_feat1, video_info_dict
elif mode == "concat":
total_spec_len = int(self.truncate / self.hop_len)
# Random Trucate len:
spec1_truncate_len = random.randint(self.min_duration * self.sr // self.hop_len, total_spec_len - self.min_duration * self.sr // self.hop_len - 1)
spec2_truncate_len = total_spec_len - spec1_truncate_len
# Sample spec clip:
spec_start1 = random.randint(0, total_spec_len - spec1_truncate_len - 1)
spec_start2 = random.randint(0, total_spec_len - spec2_truncate_len - 1)
spec_end1, spec_end2 = spec_start1 + spec1_truncate_len, spec_start2 + spec2_truncate_len
# concat spec:
return video_info_dict
def __getitem__(self, idx):
audio_name1 = self.data_list[idx]
wav_path1 = self.wav_list[idx]
video_feat_path1 = self.feat_list[idx]
# select other video:
flag = False
if random.uniform(0, 1) < -1:
flag = True
random_idx = idx
while random_idx == idx:
random_idx = random.randint(0, len(self.data_list)-1)
audio_name2 = self.data_list[random_idx]
wav_path2 = self.wav_list[random_idx]
video_feat_path2 = self.feat_list[random_idx]
# Load the Spec and Feat:
wav_raw1, video_feat1 = self.load_spec_and_feat(wav_path1, video_feat_path1)
if flag:
wav_raw2, video_feat2 = self.load_spec_and_feat(wav_path2, video_feat_path2)
video_info_dict = {'audio_name1':audio_name1, 'audio_name2': audio_name2}
mix_spec, mix_video_feat, mix_info = self.mix_audio_and_feat(video_info_dict, mode='concat')
else:
video_info_dict = {'audio_name1':audio_name1, 'audio_name2': ""}
mix_wav, mix_video_feat, mix_info = self.mix_audio_and_feat(wav_raw1=wav_raw1, video_feat1=video_feat1, video_info_dict=video_info_dict, mode='single')
data_dict = {}
data_dict['mix_wav'] = mix_wav # (131072,)
data_dict['mix_video_feat'] = mix_video_feat # (32, 512)
data_dict['mix_info_dict'] = mix_info
return data_dict
class video_codec_Dataset_Train(video_codec_Dataset):
def __init__(self, dataset_cfg):
super().__init__(split='train', **dataset_cfg)
class video_codec_Dataset_Test(video_codec_Dataset):
def __init__(self, dataset_cfg):
super().__init__(split='test', **dataset_cfg)
class video_codec_Dataset_Valid(video_codec_Dataset):
def __init__(self, dataset_cfg):
super().__init__(split='valid', **dataset_cfg)
class audio_video_spec_fullset_Audioset_Dataset(torch.utils.data.Dataset):
# Only Load audio dataset: for training Stage1: Audio Npy Dataset
def __init__(self, split, dataset1, dataset2, sr=22050, duration=10, truncate=220000,
fps=21.5, drop=0.0, fix_frames=False, hop_len=256):
super().__init__()
if split == "train":
self.split = "Train"
elif split == "valid" or split == 'test':
self.split = "Test"
# Default params:
self.min_duration = 2
self.sr = sr # 22050
self.duration = duration # 10
self.truncate = truncate # 220000
self.fps = fps
self.fix_frames = fix_frames
self.hop_len = hop_len
self.drop = drop
print("Fix Frames: {}".format(self.fix_frames))
print("Use Drop: {}".format(self.drop))
# Dataset1: (VGGSound)
assert dataset1.dataset_name == "VGGSound"
assert dataset2.dataset_name == "Audioset"
# spec_dir: spectrogram path
# feat_dir: CAVP feature path
# video_dir: video path
dataset1_spec_dir = os.path.join(dataset1.data_dir, "mel_maa2", "npy")
dataset1_feat_dir = os.path.join(dataset1.data_dir, "cavp")
split_txt_path = dataset1.split_txt_path
with open(os.path.join(split_txt_path, '{}.txt'.format(self.split)), "r") as f:
data_list1 = f.readlines()
data_list1 = list(map(lambda x: x.strip(), data_list1))
spec_list1 = list(map(lambda x: os.path.join(dataset1_spec_dir, x) + "_mel.npy", data_list1)) # spec
feat_list1 = list(map(lambda x: os.path.join(dataset1_feat_dir, x) + ".npz", data_list1)) # feat
if split == "train":
dataset2_spec_dir = os.path.join(dataset2.data_dir, "mel")
dataset2_feat_dir = os.path.join(dataset2.data_dir, "cavp_renamed")
split_txt_path = dataset2.split_txt_path
with open(os.path.join(split_txt_path, '{}.txt'.format(self.split)), "r") as f:
data_list2 = f.readlines()
data_list2 = list(map(lambda x: x.strip(), data_list2))
spec_list2 = list(map(lambda x: os.path.join(dataset2_spec_dir, f'Y{x}') + "_mel.npy", data_list2)) # spec
feat_list2 = list(map(lambda x: os.path.join(dataset2_feat_dir, x) + ".npz", data_list2)) # feat
data_list1 += data_list2
spec_list1 += spec_list2
feat_list1 += feat_list2
# Merge Data:
self.data_list = data_list1 if self.split != "Test" else data_list1[:200]
self.spec_list = spec_list1 if self.split != "Test" else spec_list1[:200]
self.feat_list = feat_list1 if self.split != "Test" else feat_list1[:200]
assert len(self.data_list) == len(self.spec_list) == len(self.feat_list)
shuffle_idx = np.random.permutation(np.arange(len(self.data_list)))
self.data_list = [self.data_list[i] for i in shuffle_idx]
self.spec_list = [self.spec_list[i] for i in shuffle_idx]
self.feat_list = [self.feat_list[i] for i in shuffle_idx]
print('Split: {} Sample Num: {}'.format(split, len(self.data_list)))
# self.check(self.spec_list)
def __len__(self):
return len(self.data_list)
def check(self, feat_list):
from tqdm import tqdm
for spec_path in tqdm(feat_list):
mel = np.load(spec_path).astype(np.float32)
if mel.shape[0] != 80:
import ipdb
ipdb.set_trace()
def load_spec_and_feat(self, spec_path, video_feat_path):
"""Load audio spec and video feat"""
spec_raw = np.load(spec_path).astype(np.float32) # channel: 1
if spec_raw.shape[0] != 80:
print(f"corrupted mel: {spec_path}", flush=True)
spec_raw = np.zeros((80, 625), dtype=np.float32) # [C, T]
p = np.random.uniform(0, 1)
if p > self.drop:
try:
video_feat = np.load(video_feat_path)['feat'].astype(np.float32)
except:
print(f"corrupted video: {video_feat_path}", flush=True)
video_feat = np.load(os.path.join(os.path.dirname(video_feat_path), 'empty_vid.npz'))['feat'].astype(np.float32)
else:
video_feat = np.load(os.path.join(os.path.dirname(video_feat_path), 'empty_vid.npz'))['feat'].astype(np.float32)
spec_len = self.sr * self.duration / self.hop_len
if spec_raw.shape[1] < spec_len:
spec_raw = np.tile(spec_raw, math.ceil(spec_len / spec_raw.shape[1]))
spec_raw = spec_raw[:, :int(spec_len)]
feat_len = self.fps * self.duration
if video_feat.shape[0] < feat_len:
video_feat = np.tile(video_feat, (math.ceil(feat_len / video_feat.shape[0]), 1))
video_feat = video_feat[:int(feat_len)]
return spec_raw, video_feat
def mix_audio_and_feat(self, spec1=None, spec2=None, video_feat1=None, video_feat2=None, video_info_dict={},
mode='single'):
""" Return Mix Spec and Mix video feat"""
if mode == "single":
# spec1:
if not self.fix_frames:
start_idx = random.randint(0, self.sr * self.duration - self.truncate - 1) # audio start
else:
start_idx = 0
start_frame = int(self.fps * start_idx / self.sr)
truncate_frame = int(self.fps * self.truncate / self.sr)
# Spec Start & Truncate:
spec_start = int(start_idx / self.hop_len)
spec_truncate = int(self.truncate / self.hop_len)
spec1 = spec1[:, spec_start: spec_start + spec_truncate]
video_feat1 = video_feat1[start_frame: start_frame + truncate_frame]
# info_dict:
video_info_dict['video_time1'] = str(start_frame) + '_' + str(
start_frame + truncate_frame) # Start frame, end frame
video_info_dict['video_time2'] = ""
return spec1, video_feat1, video_info_dict
elif mode == "concat":
total_spec_len = int(self.truncate / self.hop_len)
# Random Trucate len:
spec1_truncate_len = random.randint(self.min_duration * self.sr // self.hop_len,
total_spec_len - self.min_duration * self.sr // self.hop_len - 1)
spec2_truncate_len = total_spec_len - spec1_truncate_len
# Sample spec clip:
spec_start1 = random.randint(0, total_spec_len - spec1_truncate_len - 1)
spec_start2 = random.randint(0, total_spec_len - spec2_truncate_len - 1)
spec_end1, spec_end2 = spec_start1 + spec1_truncate_len, spec_start2 + spec2_truncate_len
# concat spec:
spec1, spec2 = spec1[:, spec_start1: spec_end1], spec2[:, spec_start2: spec_end2]
concat_audio_spec = np.concatenate([spec1, spec2], axis=1)
# Concat Video Feat:
start1_frame, truncate1_frame = int(self.fps * spec_start1 * self.hop_len / self.sr), int(
self.fps * spec1_truncate_len * self.hop_len / self.sr)
start2_frame, truncate2_frame = int(self.fps * spec_start2 * self.hop_len / self.sr), int(
self.fps * self.truncate / self.sr) - truncate1_frame
video_feat1, video_feat2 = video_feat1[start1_frame: start1_frame + truncate1_frame], video_feat2[
start2_frame: start2_frame + truncate2_frame]
concat_video_feat = np.concatenate([video_feat1, video_feat2])
video_info_dict['video_time1'] = str(start1_frame) + '_' + str(
start1_frame + truncate1_frame) # Start frame, end frame
video_info_dict['video_time2'] = str(start2_frame) + '_' + str(start2_frame + truncate2_frame)
return concat_audio_spec, concat_video_feat, video_info_dict
def __getitem__(self, idx):
audio_name1 = self.data_list[idx]
spec_npy_path1 = self.spec_list[idx]
video_feat_path1 = self.feat_list[idx]
# select other video:
flag = False
if random.uniform(0, 1) < -1:
flag = True
random_idx = idx
while random_idx == idx:
random_idx = random.randint(0, len(self.data_list) - 1)
audio_name2 = self.data_list[random_idx]
spec_npy_path2 = self.spec_list[random_idx]
video_feat_path2 = self.feat_list[random_idx]
# Load the Spec and Feat:
spec1, video_feat1 = self.load_spec_and_feat(spec_npy_path1, video_feat_path1)
if flag:
spec2, video_feat2 = self.load_spec_and_feat(spec_npy_path2, video_feat_path2)
video_info_dict = {'audio_name1': audio_name1, 'audio_name2': audio_name2}
mix_spec, mix_video_feat, mix_info = self.mix_audio_and_feat(spec1, spec2, video_feat1, video_feat2, video_info_dict, mode='concat')
else:
video_info_dict = {'audio_name1': audio_name1, 'audio_name2': ""}
mix_spec, mix_video_feat, mix_info = self.mix_audio_and_feat(spec1=spec1, video_feat1=video_feat1, video_info_dict=video_info_dict, mode='single')
# print("mix spec shape:", mix_spec.shape)
# print("mix video feat:", mix_video_feat.shape)
data_dict = {}
data_dict['mix_spec'] = mix_spec # (80, 512)
data_dict['mix_video_feat'] = mix_video_feat # (32, 512)
data_dict['mix_info_dict'] = mix_info
return data_dict
class audio_video_spec_fullset_Audioset_Train(audio_video_spec_fullset_Audioset_Dataset):
def __init__(self, dataset_cfg):
super().__init__(split='train', **dataset_cfg)
class audio_video_spec_fullset_Audioset_Valid(audio_video_spec_fullset_Audioset_Dataset):
def __init__(self, dataset_cfg):
super().__init__(split='valid', **dataset_cfg)
class audio_video_spec_fullset_Audioset_Test(audio_video_spec_fullset_Audioset_Dataset):
def __init__(self, dataset_cfg):
super().__init__(split='test', **dataset_cfg)