Spaces:
Running
on
Zero
Running
on
Zero
| from matplotlib import collections | |
| import json | |
| import os | |
| import copy | |
| import matplotlib.pyplot as plt | |
| import torch | |
| from torchvision import transforms | |
| import numpy as np | |
| from tqdm import tqdm | |
| from random import sample | |
| import torchaudio | |
| import logging | |
| import collections | |
| from glob import glob | |
| import sys | |
| import albumentations | |
| import soundfile | |
| sys.path.insert(0, '.') # nopep8 | |
| from train import instantiate_from_config | |
| from foleycrafter.models.specvqgan.data.transforms import * | |
| torchaudio.set_audio_backend("sox_io") | |
| logger = logging.getLogger(f'main.{__name__}') | |
| SR = 22050 | |
| FPS = 15 | |
| MAX_SAMPLE_ITER = 10 | |
| def non_negative(x): return int(np.round(max(0, x), 0)) | |
| def rms(x): return np.sqrt(np.mean(x**2)) | |
| def get_GH_data_identifier(video_name, start_idx, split='_'): | |
| if isinstance(start_idx, str): | |
| return video_name + split + start_idx | |
| elif isinstance(start_idx, int): | |
| return video_name + split + str(start_idx) | |
| else: | |
| raise NotImplementedError | |
| class Crop(object): | |
| def __init__(self, cropped_shape=None, random_crop=False): | |
| self.cropped_shape = cropped_shape | |
| if cropped_shape is not None: | |
| mel_num, spec_len = cropped_shape | |
| if random_crop: | |
| self.cropper = albumentations.RandomCrop | |
| else: | |
| self.cropper = albumentations.CenterCrop | |
| self.preprocessor = albumentations.Compose([self.cropper(mel_num, spec_len)]) | |
| else: | |
| self.preprocessor = lambda **kwargs: kwargs | |
| def __call__(self, item): | |
| item['image'] = self.preprocessor(image=item['image'])['image'] | |
| if 'cond_image' in item.keys(): | |
| item['cond_image'] = self.preprocessor(image=item['cond_image'])['image'] | |
| return item | |
| class CropImage(Crop): | |
| def __init__(self, *crop_args): | |
| super().__init__(*crop_args) | |
| class CropFeats(Crop): | |
| def __init__(self, *crop_args): | |
| super().__init__(*crop_args) | |
| def __call__(self, item): | |
| item['feature'] = self.preprocessor(image=item['feature'])['image'] | |
| return item | |
| class CropCoords(Crop): | |
| def __init__(self, *crop_args): | |
| super().__init__(*crop_args) | |
| def __call__(self, item): | |
| item['coord'] = self.preprocessor(image=item['coord'])['image'] | |
| return item | |
| class ResampleFrames(object): | |
| def __init__(self, feat_sample_size, times_to_repeat_after_resample=None): | |
| self.feat_sample_size = feat_sample_size | |
| self.times_to_repeat_after_resample = times_to_repeat_after_resample | |
| def __call__(self, item): | |
| feat_len = item['feature'].shape[0] | |
| ## resample | |
| assert feat_len >= self.feat_sample_size | |
| # evenly spaced points (abcdefghkl -> aoooofoooo) | |
| idx = np.linspace(0, feat_len, self.feat_sample_size, dtype=np.int, endpoint=False) | |
| # xoooo xoooo -> ooxoo ooxoo | |
| shift = feat_len // (self.feat_sample_size + 1) | |
| idx = idx + shift | |
| ## repeat after resampling (abc -> aaaabbbbcccc) | |
| if self.times_to_repeat_after_resample is not None and self.times_to_repeat_after_resample > 1: | |
| idx = np.repeat(idx, self.times_to_repeat_after_resample) | |
| item['feature'] = item['feature'][idx, :] | |
| return item | |
| class GreatestHitSpecs(torch.utils.data.Dataset): | |
| def __init__(self, split, spec_dir_path, spec_len, random_crop, mel_num, | |
| spec_crop_len, L=2.0, rand_shift=False, spec_transforms=None, splits_path='./data', | |
| meta_path='./data/info_r2plus1d_dim1024_15fps.json'): | |
| super().__init__() | |
| self.split = split | |
| self.specs_dir = spec_dir_path | |
| self.spec_transforms = spec_transforms | |
| self.splits_path = splits_path | |
| self.meta_path = meta_path | |
| self.spec_len = spec_len | |
| self.rand_shift = rand_shift | |
| self.L = L | |
| self.spec_take_first = int(math.ceil(860 * (L / 10.) / 32) * 32) | |
| self.spec_take_first = 860 if self.spec_take_first > 860 else self.spec_take_first | |
| greatesthit_meta = json.load(open(self.meta_path, 'r')) | |
| unique_classes = sorted(list(set(ht for ht in greatesthit_meta['hit_type']))) | |
| self.label2target = {label: target for target, label in enumerate(unique_classes)} | |
| self.target2label = {target: label for label, target in self.label2target.items()} | |
| self.video_idx2label = { | |
| get_GH_data_identifier(greatesthit_meta['video_name'][i], greatesthit_meta['start_idx'][i]): | |
| greatesthit_meta['hit_type'][i] for i in range(len(greatesthit_meta['video_name'])) | |
| } | |
| self.available_video_hit = list(self.video_idx2label.keys()) | |
| self.video_idx2path = { | |
| vh: os.path.join(self.specs_dir, | |
| vh.replace('_', '_denoised_') + '_' + self.video_idx2label[vh].replace(' ', '_') +'_mel.npy') | |
| for vh in self.available_video_hit | |
| } | |
| self.video_idx2idx = { | |
| get_GH_data_identifier(greatesthit_meta['video_name'][i], greatesthit_meta['start_idx'][i]): | |
| i for i in range(len(greatesthit_meta['video_name'])) | |
| } | |
| split_clip_ids_path = os.path.join(splits_path, f'greatesthit_{split}.json') | |
| if not os.path.exists(split_clip_ids_path): | |
| raise NotImplementedError() | |
| clip_video_hit = json.load(open(split_clip_ids_path, 'r')) | |
| self.dataset = clip_video_hit | |
| spec_crop_len = self.spec_take_first if self.spec_take_first <= spec_crop_len else spec_crop_len | |
| self.spec_transforms = transforms.Compose([ | |
| CropImage([mel_num, spec_crop_len], random_crop), | |
| # transforms.RandomApply([FrequencyMasking(freq_mask_param=20)], p=0), | |
| # transforms.RandomApply([TimeMasking(time_mask_param=int(32 * self.L))], p=0) | |
| ]) | |
| self.video2indexes = {} | |
| for video_idx in self.dataset: | |
| video, start_idx = video_idx.split('_') | |
| if video not in self.video2indexes.keys(): | |
| self.video2indexes[video] = [] | |
| self.video2indexes[video].append(start_idx) | |
| for video in self.video2indexes.keys(): | |
| if len(self.video2indexes[video]) == 1: # given video contains only one hit | |
| self.dataset.remove( | |
| get_GH_data_identifier(video, self.video2indexes[video][0]) | |
| ) | |
| def __len__(self): | |
| return len(self.dataset) | |
| def __getitem__(self, idx): | |
| item = {} | |
| video_idx = self.dataset[idx] | |
| spec_path = self.video_idx2path[video_idx] | |
| spec = np.load(spec_path) # (80, 860) | |
| if self.rand_shift: | |
| shift = random.uniform(0, 0.5) | |
| spec_shift = int(shift * spec.shape[1] // 10) | |
| # Since only the first second is used | |
| spec = np.roll(spec, -spec_shift, 1) | |
| # concat spec outside dataload | |
| item['image'] = 2 * spec - 1 # (80, 860) | |
| item['image'] = item['image'][:, :self.spec_take_first] | |
| item['file_path'] = spec_path | |
| item['label'] = self.video_idx2label[video_idx] | |
| item['target'] = self.label2target[item['label']] | |
| if self.spec_transforms is not None: | |
| item = self.spec_transforms(item) | |
| return item | |
| class GreatestHitSpecsTrain(GreatestHitSpecs): | |
| def __init__(self, specs_dataset_cfg): | |
| super().__init__('train', **specs_dataset_cfg) | |
| class GreatestHitSpecsValidation(GreatestHitSpecs): | |
| def __init__(self, specs_dataset_cfg): | |
| super().__init__('val', **specs_dataset_cfg) | |
| class GreatestHitSpecsTest(GreatestHitSpecs): | |
| def __init__(self, specs_dataset_cfg): | |
| super().__init__('test', **specs_dataset_cfg) | |
| class GreatestHitWave(torch.utils.data.Dataset): | |
| def __init__(self, split, wav_dir, random_crop, mel_num, spec_crop_len, spec_len, | |
| L=2.0, splits_path='./data', rand_shift=True, | |
| data_path='data/greatesthit/greatesthit-process-resized'): | |
| super().__init__() | |
| self.split = split | |
| self.wav_dir = wav_dir | |
| self.splits_path = splits_path | |
| self.data_path = data_path | |
| self.L = L | |
| self.rand_shift = rand_shift | |
| split_clip_ids_path = os.path.join(splits_path, f'greatesthit_{split}.json') | |
| if not os.path.exists(split_clip_ids_path): | |
| raise NotImplementedError() | |
| clip_video_hit = json.load(open(split_clip_ids_path, 'r')) | |
| video_name = list(set([vidx.split('_')[0] for vidx in clip_video_hit])) | |
| self.video_frame_cnt = {v: len(os.listdir(os.path.join(self.data_path, v, 'frames'))) // 2 for v in video_name} | |
| self.left_over = int(FPS * L + 1) | |
| self.video_audio_path = {v: os.path.join(self.data_path, v, f'audio/{v}_denoised_resampled.wav') for v in video_name} | |
| self.dataset = clip_video_hit | |
| self.video2indexes = {} | |
| for video_idx in self.dataset: | |
| video, start_idx = video_idx.split('_') | |
| if video not in self.video2indexes.keys(): | |
| self.video2indexes[video] = [] | |
| self.video2indexes[video].append(start_idx) | |
| for video in self.video2indexes.keys(): | |
| if len(self.video2indexes[video]) == 1: # given video contains only one hit | |
| self.dataset.remove( | |
| get_GH_data_identifier(video, self.video2indexes[video][0]) | |
| ) | |
| self.wav_transforms = transforms.Compose([ | |
| MakeMono(), | |
| Padding(target_len=int(SR * self.L)), | |
| ]) | |
| def __len__(self): | |
| return len(self.dataset) | |
| def __getitem__(self, idx): | |
| item = {} | |
| video_idx = self.dataset[idx] | |
| video, start_idx = video_idx.split('_') | |
| start_idx = int(start_idx) | |
| if self.rand_shift: | |
| shift = int(random.uniform(-0.5, 0.5) * SR) | |
| start_idx = non_negative(start_idx + shift) | |
| wave_path = self.video_audio_path[video] | |
| wav, sr = soundfile.read(wave_path, frames=int(SR * self.L), start=start_idx) | |
| assert sr == SR | |
| wav = self.wav_transforms(wav) | |
| item['image'] = wav # (44100,) | |
| # item['wav'] = wav | |
| item['file_path_wav_'] = wave_path | |
| item['label'] = 'None' | |
| item['target'] = 'None' | |
| return item | |
| class GreatestHitWaveTrain(GreatestHitWave): | |
| def __init__(self, specs_dataset_cfg): | |
| super().__init__('train', **specs_dataset_cfg) | |
| class GreatestHitWaveValidation(GreatestHitWave): | |
| def __init__(self, specs_dataset_cfg): | |
| super().__init__('val', **specs_dataset_cfg) | |
| class GreatestHitWaveTest(GreatestHitWave): | |
| def __init__(self, specs_dataset_cfg): | |
| super().__init__('test', **specs_dataset_cfg) | |
| class CondGreatestHitSpecsCondOnImage(torch.utils.data.Dataset): | |
| def __init__(self, split, specs_dir, spec_len, feat_len, feat_depth, feat_crop_len, random_crop, mel_num, spec_crop_len, | |
| vqgan_L=10.0, L=1.0, rand_shift=False, spec_transforms=None, frame_transforms=None, splits_path='./data', | |
| meta_path='./data/info_r2plus1d_dim1024_15fps.json', frame_path='data/greatesthit/greatesthit_processed', | |
| p_outside_cond=0., p_audio_aug=0.5): | |
| super().__init__() | |
| self.split = split | |
| self.specs_dir = specs_dir | |
| self.spec_transforms = spec_transforms | |
| self.frame_transforms = frame_transforms | |
| self.splits_path = splits_path | |
| self.meta_path = meta_path | |
| self.frame_path = frame_path | |
| self.feat_len = feat_len | |
| self.feat_depth = feat_depth | |
| self.feat_crop_len = feat_crop_len | |
| self.spec_len = spec_len | |
| self.rand_shift = rand_shift | |
| self.L = L | |
| self.spec_take_first = int(math.ceil(860 * (vqgan_L / 10.) / 32) * 32) | |
| self.spec_take_first = 860 if self.spec_take_first > 860 else self.spec_take_first | |
| self.p_outside_cond = torch.tensor(p_outside_cond) | |
| greatesthit_meta = json.load(open(self.meta_path, 'r')) | |
| unique_classes = sorted(list(set(ht for ht in greatesthit_meta['hit_type']))) | |
| self.label2target = {label: target for target, label in enumerate(unique_classes)} | |
| self.target2label = {target: label for label, target in self.label2target.items()} | |
| self.video_idx2label = { | |
| get_GH_data_identifier(greatesthit_meta['video_name'][i], greatesthit_meta['start_idx'][i]): | |
| greatesthit_meta['hit_type'][i] for i in range(len(greatesthit_meta['video_name'])) | |
| } | |
| self.available_video_hit = list(self.video_idx2label.keys()) | |
| self.video_idx2path = { | |
| vh: os.path.join(self.specs_dir, | |
| vh.replace('_', '_denoised_') + '_' + self.video_idx2label[vh].replace(' ', '_') +'_mel.npy') | |
| for vh in self.available_video_hit | |
| } | |
| for value in self.video_idx2path.values(): | |
| assert os.path.exists(value) | |
| self.video_idx2idx = { | |
| get_GH_data_identifier(greatesthit_meta['video_name'][i], greatesthit_meta['start_idx'][i]): | |
| i for i in range(len(greatesthit_meta['video_name'])) | |
| } | |
| split_clip_ids_path = os.path.join(splits_path, f'greatesthit_{split}.json') | |
| if not os.path.exists(split_clip_ids_path): | |
| self.make_split_files() | |
| clip_video_hit = json.load(open(split_clip_ids_path, 'r')) | |
| self.dataset = clip_video_hit | |
| spec_crop_len = self.spec_take_first if self.spec_take_first <= spec_crop_len else spec_crop_len | |
| self.spec_transforms = transforms.Compose([ | |
| CropImage([mel_num, spec_crop_len], random_crop), | |
| # transforms.RandomApply([FrequencyMasking(freq_mask_param=20)], p=p_audio_aug), | |
| # transforms.RandomApply([TimeMasking(time_mask_param=int(32 * self.L))], p=p_audio_aug) | |
| ]) | |
| if self.frame_transforms == None: | |
| self.frame_transforms = transforms.Compose([ | |
| Resize3D(128), | |
| RandomResizedCrop3D(112, scale=(0.5, 1.0)), | |
| RandomHorizontalFlip3D(), | |
| ColorJitter3D(brightness=0.1, saturation=0.1), | |
| ToTensor3D(), | |
| Normalize3D(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]), | |
| ]) | |
| self.video2indexes = {} | |
| for video_idx in self.dataset: | |
| video, start_idx = video_idx.split('_') | |
| if video not in self.video2indexes.keys(): | |
| self.video2indexes[video] = [] | |
| self.video2indexes[video].append(start_idx) | |
| for video in self.video2indexes.keys(): | |
| if len(self.video2indexes[video]) == 1: # given video contains only one hit | |
| self.dataset.remove( | |
| get_GH_data_identifier(video, self.video2indexes[video][0]) | |
| ) | |
| clip_classes = [self.label2target[self.video_idx2label[vh]] for vh in clip_video_hit] | |
| class2count = collections.Counter(clip_classes) | |
| self.class_counts = torch.tensor([class2count[cls] for cls in range(len(class2count))]) | |
| if self.L != 1.0: | |
| print(split, L) | |
| self.validate_data() | |
| self.video2indexes = {} | |
| for video_idx in self.dataset: | |
| video, start_idx = video_idx.split('_') | |
| if video not in self.video2indexes.keys(): | |
| self.video2indexes[video] = [] | |
| self.video2indexes[video].append(start_idx) | |
| def __len__(self): | |
| return len(self.dataset) | |
| def __getitem__(self, idx): | |
| item = {} | |
| try: | |
| video_idx = self.dataset[idx] | |
| spec_path = self.video_idx2path[video_idx] | |
| spec = np.load(spec_path) # (80, 860) | |
| video, start_idx = video_idx.split('_') | |
| frame_path = os.path.join(self.frame_path, video, 'frames') | |
| start_frame_idx = non_negative(FPS * int(start_idx)/SR) | |
| end_frame_idx = non_negative(start_frame_idx + FPS * self.L) | |
| if self.rand_shift: | |
| shift = random.uniform(0, 0.5) | |
| spec_shift = int(shift * spec.shape[1] // 10) | |
| # Since only the first second is used | |
| spec = np.roll(spec, -spec_shift, 1) | |
| start_frame_idx += int(FPS * shift) | |
| end_frame_idx += int(FPS * shift) | |
| frames = [Image.open(os.path.join( | |
| frame_path, f'frame{i+1:0>6d}.jpg')).convert('RGB') for i in | |
| range(start_frame_idx, end_frame_idx)] | |
| # Sample condition | |
| if torch.all(torch.bernoulli(self.p_outside_cond) == 1.): | |
| # Sample condition from outside video | |
| all_idx = set(list(range(len(self.dataset)))) | |
| all_idx.remove(idx) | |
| cond_video_idx = self.dataset[sample(all_idx, k=1)[0]] | |
| cond_video, cond_start_idx = cond_video_idx.split('_') | |
| else: | |
| cond_video = video | |
| video_hits_idx = copy.copy(self.video2indexes[video]) | |
| video_hits_idx.remove(start_idx) | |
| cond_start_idx = sample(video_hits_idx, k=1)[0] | |
| cond_video_idx = get_GH_data_identifier(cond_video, cond_start_idx) | |
| cond_spec_path = self.video_idx2path[cond_video_idx] | |
| cond_spec = np.load(cond_spec_path) # (80, 860) | |
| cond_video, cond_start_idx = cond_video_idx.split('_') | |
| cond_frame_path = os.path.join(self.frame_path, cond_video, 'frames') | |
| cond_start_frame_idx = non_negative(FPS * int(cond_start_idx)/SR) | |
| cond_end_frame_idx = non_negative(cond_start_frame_idx + FPS * self.L) | |
| if self.rand_shift: | |
| cond_shift = random.uniform(0, 0.5) | |
| cond_spec_shift = int(cond_shift * cond_spec.shape[1] // 10) | |
| # Since only the first second is used | |
| cond_spec = np.roll(cond_spec, -cond_spec_shift, 1) | |
| cond_start_frame_idx += int(FPS * cond_shift) | |
| cond_end_frame_idx += int(FPS * cond_shift) | |
| cond_frames = [Image.open(os.path.join( | |
| cond_frame_path, f'frame{i+1:0>6d}.jpg')).convert('RGB') for i in | |
| range(cond_start_frame_idx, cond_end_frame_idx)] | |
| # concat spec outside dataload | |
| item['image'] = 2 * spec - 1 # (80, 860) | |
| item['cond_image'] = 2 * cond_spec - 1 # (80, 860) | |
| item['image'] = item['image'][:, :self.spec_take_first] | |
| item['cond_image'] = item['cond_image'][:, :self.spec_take_first] | |
| item['file_path_specs_'] = spec_path | |
| item['file_path_cond_specs_'] = cond_spec_path | |
| if self.frame_transforms is not None: | |
| cond_frames = self.frame_transforms(cond_frames) | |
| frames = self.frame_transforms(frames) | |
| item['feature'] = np.stack(cond_frames + frames, axis=0) # (30 * L, 112, 112, 3) | |
| item['file_path_feats_'] = (frame_path, start_frame_idx) | |
| item['file_path_cond_feats_'] = (cond_frame_path, cond_start_frame_idx) | |
| item['label'] = self.video_idx2label[video_idx] | |
| item['target'] = self.label2target[item['label']] | |
| if self.spec_transforms is not None: | |
| item = self.spec_transforms(item) | |
| except Exception: | |
| print(sys.exc_info()[2]) | |
| print('!!!!!!!!!!!!!!!!!!!!', video_idx, cond_video_idx) | |
| print('!!!!!!!!!!!!!!!!!!!!', end_frame_idx, cond_end_frame_idx) | |
| exit(1) | |
| return item | |
| def validate_data(self): | |
| original_len = len(self.dataset) | |
| valid_dataset = [] | |
| for video_idx in tqdm(self.dataset): | |
| video, start_idx = video_idx.split('_') | |
| frame_path = os.path.join(self.frame_path, video, 'frames') | |
| start_frame_idx = non_negative(FPS * int(start_idx)/SR) | |
| end_frame_idx = non_negative(start_frame_idx + FPS * (self.L + 0.6)) | |
| if os.path.exists(os.path.join(frame_path, f'frame{end_frame_idx:0>6d}.jpg')): | |
| valid_dataset.append(video_idx) | |
| else: | |
| self.video2indexes[video].remove(start_idx) | |
| for video_idx in valid_dataset: | |
| video, start_idx = video_idx.split('_') | |
| if len(self.video2indexes[video]) == 1: | |
| valid_dataset.remove(video_idx) | |
| if original_len != len(valid_dataset): | |
| print(f'Validated dataset with enough frames: {len(valid_dataset)}') | |
| self.dataset = valid_dataset | |
| split_clip_ids_path = os.path.join(self.splits_path, f'greatesthit_{self.split}_{self.L:.2f}.json') | |
| if not os.path.exists(split_clip_ids_path): | |
| with open(split_clip_ids_path, 'w') as f: | |
| json.dump(valid_dataset, f) | |
| def make_split_files(self, ratio=[0.85, 0.1, 0.05]): | |
| random.seed(1337) | |
| print(f'The split files do not exist @ {self.splits_path}. Calculating the new ones.') | |
| # The downloaded videos (some went missing on YouTube and no longer available) | |
| available_mel_paths = set(glob(os.path.join(self.specs_dir, '*_mel.npy'))) | |
| self.available_video_hit = [vh for vh in self.available_video_hit if self.video_idx2path[vh] in available_mel_paths] | |
| all_video = list(self.video2indexes.keys()) | |
| print(f'The number of clips available after download: {len(self.available_video_hit)}') | |
| print(f'The number of videos available after download: {len(all_video)}') | |
| available_idx = list(range(len(all_video))) | |
| random.shuffle(available_idx) | |
| assert sum(ratio) == 1. | |
| cut_train = int(ratio[0] * len(all_video)) | |
| cut_test = cut_train + int(ratio[1] * len(all_video)) | |
| train_idx = available_idx[:cut_train] | |
| test_idx = available_idx[cut_train:cut_test] | |
| valid_idx = available_idx[cut_test:] | |
| train_video = [all_video[i] for i in train_idx] | |
| test_video = [all_video[i] for i in test_idx] | |
| valid_video = [all_video[i] for i in valid_idx] | |
| train_video_hit = [] | |
| for v in train_video: | |
| train_video_hit += [get_GH_data_identifier(v, hit_idx) for hit_idx in self.video2indexes[v]] | |
| test_video_hit = [] | |
| for v in test_video: | |
| test_video_hit += [get_GH_data_identifier(v, hit_idx) for hit_idx in self.video2indexes[v]] | |
| valid_video_hit = [] | |
| for v in valid_video: | |
| valid_video_hit += [get_GH_data_identifier(v, hit_idx) for hit_idx in self.video2indexes[v]] | |
| # mix train and valid for better validation loss | |
| mixed = train_video_hit + valid_video_hit | |
| random.shuffle(mixed) | |
| split = int(len(mixed) * ratio[0] / (ratio[0] + ratio[2])) | |
| train_video_hit = mixed[:split] | |
| valid_video_hit = mixed[split:] | |
| with open(os.path.join(self.splits_path, 'greatesthit_train.json'), 'w') as train_file,\ | |
| open(os.path.join(self.splits_path, 'greatesthit_test.json'), 'w') as test_file,\ | |
| open(os.path.join(self.splits_path, 'greatesthit_valid.json'), 'w') as valid_file: | |
| json.dump(train_video_hit, train_file) | |
| json.dump(test_video_hit, test_file) | |
| json.dump(valid_video_hit, valid_file) | |
| print(f'Put {len(train_idx)} clips to the train set and saved it to ./data/greatesthit_train.json') | |
| print(f'Put {len(test_idx)} clips to the test set and saved it to ./data/greatesthit_test.json') | |
| print(f'Put {len(valid_idx)} clips to the valid set and saved it to ./data/greatesthit_valid.json') | |
| class CondGreatestHitSpecsCondOnImageTrain(CondGreatestHitSpecsCondOnImage): | |
| def __init__(self, dataset_cfg): | |
| train_transforms = transforms.Compose([ | |
| Resize3D(256), | |
| RandomResizedCrop3D(224, scale=(0.5, 1.0)), | |
| RandomHorizontalFlip3D(), | |
| ColorJitter3D(brightness=0.1, saturation=0.1), | |
| ToTensor3D(), | |
| Normalize3D(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]), | |
| ]) | |
| super().__init__('train', frame_transforms=train_transforms, **dataset_cfg) | |
| class CondGreatestHitSpecsCondOnImageValidation(CondGreatestHitSpecsCondOnImage): | |
| def __init__(self, dataset_cfg): | |
| valid_transforms = transforms.Compose([ | |
| Resize3D(256), | |
| CenterCrop3D(224), | |
| ToTensor3D(), | |
| Normalize3D(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]), | |
| ]) | |
| super().__init__('val', frame_transforms=valid_transforms, **dataset_cfg) | |
| class CondGreatestHitSpecsCondOnImageTest(CondGreatestHitSpecsCondOnImage): | |
| def __init__(self, dataset_cfg): | |
| test_transforms = transforms.Compose([ | |
| Resize3D(256), | |
| CenterCrop3D(224), | |
| ToTensor3D(), | |
| Normalize3D(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]), | |
| ]) | |
| super().__init__('test', frame_transforms=test_transforms, **dataset_cfg) | |
| class CondGreatestHitWaveCondOnImage(torch.utils.data.Dataset): | |
| def __init__(self, split, wav_dir, spec_len, random_crop, mel_num, spec_crop_len, | |
| L=2.0, frame_transforms=None, splits_path='./data', | |
| data_path='data/greatesthit/greatesthit-process-resized', | |
| p_outside_cond=0., p_audio_aug=0.5, rand_shift=True): | |
| super().__init__() | |
| self.split = split | |
| self.wav_dir = wav_dir | |
| self.frame_transforms = frame_transforms | |
| self.splits_path = splits_path | |
| self.data_path = data_path | |
| self.spec_len = spec_len | |
| self.L = L | |
| self.rand_shift = rand_shift | |
| self.p_outside_cond = torch.tensor(p_outside_cond) | |
| split_clip_ids_path = os.path.join(splits_path, f'greatesthit_{split}.json') | |
| if not os.path.exists(split_clip_ids_path): | |
| raise NotImplementedError() | |
| clip_video_hit = json.load(open(split_clip_ids_path, 'r')) | |
| video_name = list(set([vidx.split('_')[0] for vidx in clip_video_hit])) | |
| self.video_frame_cnt = {v: len(os.listdir(os.path.join(self.data_path, v, 'frames')))//2 for v in video_name} | |
| self.left_over = int(FPS * L + 1) | |
| self.video_audio_path = {v: os.path.join(self.data_path, v, f'audio/{v}_denoised_resampled.wav') for v in video_name} | |
| self.dataset = clip_video_hit | |
| self.video2indexes = {} | |
| for video_idx in self.dataset: | |
| video, start_idx = video_idx.split('_') | |
| if video not in self.video2indexes.keys(): | |
| self.video2indexes[video] = [] | |
| self.video2indexes[video].append(start_idx) | |
| for video in self.video2indexes.keys(): | |
| if len(self.video2indexes[video]) == 1: # given video contains only one hit | |
| self.dataset.remove( | |
| get_GH_data_identifier(video, self.video2indexes[video][0]) | |
| ) | |
| self.wav_transforms = transforms.Compose([ | |
| MakeMono(), | |
| Padding(target_len=int(SR * self.L)), | |
| ]) | |
| if self.frame_transforms == None: | |
| self.frame_transforms = transforms.Compose([ | |
| Resize3D(256), | |
| RandomResizedCrop3D(224, scale=(0.5, 1.0)), | |
| RandomHorizontalFlip3D(), | |
| ColorJitter3D(brightness=0.1, saturation=0.1), | |
| ToTensor3D(), | |
| Normalize3D(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]), | |
| ]) | |
| def __len__(self): | |
| return len(self.dataset) | |
| def __getitem__(self, idx): | |
| item = {} | |
| video_idx = self.dataset[idx] | |
| video, start_idx = video_idx.split('_') | |
| start_idx = int(start_idx) | |
| frame_path = os.path.join(self.data_path, video, 'frames') | |
| start_frame_idx = non_negative(FPS * int(start_idx)/SR) | |
| if self.rand_shift: | |
| shift = random.uniform(-0.5, 0.5) | |
| start_frame_idx = non_negative(start_frame_idx + int(FPS * shift)) | |
| start_idx = non_negative(start_idx + int(SR * shift)) | |
| if start_frame_idx > self.video_frame_cnt[video] - self.left_over: | |
| start_frame_idx = self.video_frame_cnt[video] - self.left_over | |
| start_idx = non_negative(SR * (start_frame_idx / FPS)) | |
| end_frame_idx = non_negative(start_frame_idx + FPS * self.L) | |
| # target | |
| wave_path = self.video_audio_path[video] | |
| frames = [Image.open(os.path.join( | |
| frame_path, f'frame{i+1:0>6d}')).convert('RGB') for i in | |
| range(start_frame_idx, end_frame_idx)] | |
| wav, sr = soundfile.read(wave_path, frames=int(SR * self.L), start=start_idx) | |
| assert sr == SR | |
| wav = self.wav_transforms(wav) | |
| # cond | |
| if torch.all(torch.bernoulli(self.p_outside_cond) == 1.): | |
| all_idx = set(list(range(len(self.dataset)))) | |
| all_idx.remove(idx) | |
| cond_video_idx = self.dataset[sample(all_idx, k=1)[0]] | |
| cond_video, cond_start_idx = cond_video_idx.split('_') | |
| else: | |
| cond_video = video | |
| video_hits_idx = copy.copy(self.video2indexes[video]) | |
| if str(start_idx) in video_hits_idx: | |
| video_hits_idx.remove(str(start_idx)) | |
| cond_start_idx = sample(video_hits_idx, k=1)[0] | |
| cond_video_idx = get_GH_data_identifier(cond_video, cond_start_idx) | |
| cond_video, cond_start_idx = cond_video_idx.split('_') | |
| cond_start_idx = int(cond_start_idx) | |
| cond_frame_path = os.path.join(self.data_path, cond_video, 'frames') | |
| cond_start_frame_idx = non_negative(FPS * int(cond_start_idx)/SR) | |
| cond_wave_path = self.video_audio_path[cond_video] | |
| if self.rand_shift: | |
| cond_shift = random.uniform(-0.5, 0.5) | |
| cond_start_frame_idx = non_negative(cond_start_frame_idx + int(FPS * cond_shift)) | |
| cond_start_idx = non_negative(cond_start_idx + int(shift * SR)) | |
| if cond_start_frame_idx > self.video_frame_cnt[cond_video] - self.left_over: | |
| cond_start_frame_idx = self.video_frame_cnt[cond_video] - self.left_over | |
| cond_start_idx = non_negative(SR * (cond_start_frame_idx / FPS)) | |
| cond_end_frame_idx = non_negative(cond_start_frame_idx + FPS * self.L) | |
| cond_frames = [Image.open(os.path.join( | |
| cond_frame_path, f'frame{i+1:0>6d}')).convert('RGB') for i in | |
| range(cond_start_frame_idx, cond_end_frame_idx)] | |
| cond_wav, _ = soundfile.read(cond_wave_path, frames=int(SR * self.L), start=cond_start_idx) | |
| cond_wav = self.wav_transforms(cond_wav) | |
| item['image'] = wav # (44100,) | |
| item['cond_image'] = cond_wav # (44100,) | |
| item['file_path_wav_'] = wave_path | |
| item['file_path_cond_wav_'] = cond_wave_path | |
| if self.frame_transforms is not None: | |
| cond_frames = self.frame_transforms(cond_frames) | |
| frames = self.frame_transforms(frames) | |
| item['feature'] = np.stack(cond_frames + frames, axis=0) # (30 * L, 112, 112, 3) | |
| item['file_path_feats_'] = (frame_path, start_idx) | |
| item['file_path_cond_feats_'] = (cond_frame_path, cond_start_idx) | |
| item['label'] = 'None' | |
| item['target'] = 'None' | |
| return item | |
| def validate_data(self): | |
| raise NotImplementedError() | |
| def make_split_files(self, ratio=[0.85, 0.1, 0.05]): | |
| random.seed(1337) | |
| print(f'The split files do not exist @ {self.splits_path}. Calculating the new ones.') | |
| all_video = sorted(os.listdir(self.data_path)) | |
| print(f'The number of videos available after download: {len(all_video)}') | |
| available_idx = list(range(len(all_video))) | |
| random.shuffle(available_idx) | |
| assert sum(ratio) == 1. | |
| cut_train = int(ratio[0] * len(all_video)) | |
| cut_test = cut_train + int(ratio[1] * len(all_video)) | |
| train_idx = available_idx[:cut_train] | |
| test_idx = available_idx[cut_train:cut_test] | |
| valid_idx = available_idx[cut_test:] | |
| train_video = [all_video[i] for i in train_idx] | |
| test_video = [all_video[i] for i in test_idx] | |
| valid_video = [all_video[i] for i in valid_idx] | |
| with open(os.path.join(self.splits_path, 'greatesthit_video_train.json'), 'w') as train_file,\ | |
| open(os.path.join(self.splits_path, 'greatesthit_video_test.json'), 'w') as test_file,\ | |
| open(os.path.join(self.splits_path, 'greatesthit_video_valid.json'), 'w') as valid_file: | |
| json.dump(train_video, train_file) | |
| json.dump(test_video, test_file) | |
| json.dump(valid_video, valid_file) | |
| print(f'Put {len(train_idx)} videos to the train set and saved it to ./data/greatesthit_video_train.json') | |
| print(f'Put {len(test_idx)} videos to the test set and saved it to ./data/greatesthit_video_test.json') | |
| print(f'Put {len(valid_idx)} videos to the valid set and saved it to ./data/greatesthit_video_valid.json') | |
| class CondGreatestHitWaveCondOnImageTrain(CondGreatestHitWaveCondOnImage): | |
| def __init__(self, dataset_cfg): | |
| train_transforms = transforms.Compose([ | |
| Resize3D(128), | |
| RandomResizedCrop3D(112, scale=(0.5, 1.0)), | |
| RandomHorizontalFlip3D(), | |
| ColorJitter3D(brightness=0.4, saturation=0.4, contrast=0.2, hue=0.1), | |
| ToTensor3D(), | |
| Normalize3D(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]), | |
| ]) | |
| super().__init__('train', frame_transforms=train_transforms, **dataset_cfg) | |
| class CondGreatestHitWaveCondOnImageValidation(CondGreatestHitWaveCondOnImage): | |
| def __init__(self, dataset_cfg): | |
| valid_transforms = transforms.Compose([ | |
| Resize3D(128), | |
| CenterCrop3D(112), | |
| ToTensor3D(), | |
| Normalize3D(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]), | |
| ]) | |
| super().__init__('val', frame_transforms=valid_transforms, **dataset_cfg) | |
| class CondGreatestHitWaveCondOnImageTest(CondGreatestHitWaveCondOnImage): | |
| def __init__(self, dataset_cfg): | |
| test_transforms = transforms.Compose([ | |
| Resize3D(128), | |
| CenterCrop3D(112), | |
| ToTensor3D(), | |
| Normalize3D(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]), | |
| ]) | |
| super().__init__('test', frame_transforms=test_transforms, **dataset_cfg) | |
| class GreatestHitWaveCondOnImage(torch.utils.data.Dataset): | |
| def __init__(self, split, wav_dir, spec_len, random_crop, mel_num, spec_crop_len, | |
| L=2.0, frame_transforms=None, splits_path='./data', | |
| data_path='data/greatesthit/greatesthit-process-resized', | |
| p_outside_cond=0., p_audio_aug=0.5, rand_shift=True): | |
| super().__init__() | |
| self.split = split | |
| self.wav_dir = wav_dir | |
| self.frame_transforms = frame_transforms | |
| self.splits_path = splits_path | |
| self.data_path = data_path | |
| self.spec_len = spec_len | |
| self.L = L | |
| self.rand_shift = rand_shift | |
| self.p_outside_cond = torch.tensor(p_outside_cond) | |
| split_clip_ids_path = os.path.join(splits_path, f'greatesthit_{split}.json') | |
| if not os.path.exists(split_clip_ids_path): | |
| raise NotImplementedError() | |
| clip_video_hit = json.load(open(split_clip_ids_path, 'r')) | |
| video_name = list(set([vidx.split('_')[0] for vidx in clip_video_hit])) | |
| self.video_frame_cnt = {v: len(os.listdir(os.path.join(self.data_path, v, 'frames')))//2 for v in video_name} | |
| self.left_over = int(FPS * L + 1) | |
| self.video_audio_path = {v: os.path.join(self.data_path, v, f'audio/{v}_denoised_resampled.wav') for v in video_name} | |
| self.dataset = clip_video_hit | |
| self.video2indexes = {} | |
| for video_idx in self.dataset: | |
| video, start_idx = video_idx.split('_') | |
| if video not in self.video2indexes.keys(): | |
| self.video2indexes[video] = [] | |
| self.video2indexes[video].append(start_idx) | |
| for video in self.video2indexes.keys(): | |
| if len(self.video2indexes[video]) == 1: # given video contains only one hit | |
| self.dataset.remove( | |
| get_GH_data_identifier(video, self.video2indexes[video][0]) | |
| ) | |
| self.wav_transforms = transforms.Compose([ | |
| MakeMono(), | |
| Padding(target_len=int(SR * self.L)), | |
| ]) | |
| if self.frame_transforms == None: | |
| self.frame_transforms = transforms.Compose([ | |
| Resize3D(256), | |
| RandomResizedCrop3D(224, scale=(0.5, 1.0)), | |
| RandomHorizontalFlip3D(), | |
| ColorJitter3D(brightness=0.1, saturation=0.1), | |
| ToTensor3D(), | |
| Normalize3D(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]), | |
| ]) | |
| def __len__(self): | |
| return len(self.dataset) | |
| def __getitem__(self, idx): | |
| item = {} | |
| video_idx = self.dataset[idx] | |
| video, start_idx = video_idx.split('_') | |
| start_idx = int(start_idx) | |
| frame_path = os.path.join(self.data_path, video, 'frames') | |
| start_frame_idx = non_negative(FPS * int(start_idx)/SR) | |
| if self.rand_shift: | |
| shift = random.uniform(-0.5, 0.5) | |
| start_frame_idx = non_negative(start_frame_idx + int(FPS * shift)) | |
| start_idx = non_negative(start_idx + int(SR * shift)) | |
| if start_frame_idx > self.video_frame_cnt[video] - self.left_over: | |
| start_frame_idx = self.video_frame_cnt[video] - self.left_over | |
| start_idx = non_negative(SR * (start_frame_idx / FPS)) | |
| end_frame_idx = non_negative(start_frame_idx + FPS * self.L) | |
| # target | |
| wave_path = self.video_audio_path[video] | |
| frames = [Image.open(os.path.join( | |
| frame_path, f'frame{i+1:0>6d}')).convert('RGB') for i in | |
| range(start_frame_idx, end_frame_idx)] | |
| wav, sr = soundfile.read(wave_path, frames=int(SR * self.L), start=start_idx) | |
| assert sr == SR | |
| wav = self.wav_transforms(wav) | |
| item['image'] = wav # (44100,) | |
| item['file_path_wav_'] = wave_path | |
| if self.frame_transforms is not None: | |
| frames = self.frame_transforms(frames) | |
| item['feature'] = torch.stack(frames, dim=0) # (15 * L, 112, 112, 3) | |
| item['file_path_feats_'] = (frame_path, start_idx) | |
| item['label'] = 'None' | |
| item['target'] = 'None' | |
| return item | |
| def validate_data(self): | |
| raise NotImplementedError() | |
| def make_split_files(self, ratio=[0.85, 0.1, 0.05]): | |
| random.seed(1337) | |
| print(f'The split files do not exist @ {self.splits_path}. Calculating the new ones.') | |
| all_video = sorted(os.listdir(self.data_path)) | |
| print(f'The number of videos available after download: {len(all_video)}') | |
| available_idx = list(range(len(all_video))) | |
| random.shuffle(available_idx) | |
| assert sum(ratio) == 1. | |
| cut_train = int(ratio[0] * len(all_video)) | |
| cut_test = cut_train + int(ratio[1] * len(all_video)) | |
| train_idx = available_idx[:cut_train] | |
| test_idx = available_idx[cut_train:cut_test] | |
| valid_idx = available_idx[cut_test:] | |
| train_video = [all_video[i] for i in train_idx] | |
| test_video = [all_video[i] for i in test_idx] | |
| valid_video = [all_video[i] for i in valid_idx] | |
| with open(os.path.join(self.splits_path, 'greatesthit_video_train.json'), 'w') as train_file,\ | |
| open(os.path.join(self.splits_path, 'greatesthit_video_test.json'), 'w') as test_file,\ | |
| open(os.path.join(self.splits_path, 'greatesthit_video_valid.json'), 'w') as valid_file: | |
| json.dump(train_video, train_file) | |
| json.dump(test_video, test_file) | |
| json.dump(valid_video, valid_file) | |
| print(f'Put {len(train_idx)} videos to the train set and saved it to ./data/greatesthit_video_train.json') | |
| print(f'Put {len(test_idx)} videos to the test set and saved it to ./data/greatesthit_video_test.json') | |
| print(f'Put {len(valid_idx)} videos to the valid set and saved it to ./data/greatesthit_video_valid.json') | |
| class GreatestHitWaveCondOnImageTrain(GreatestHitWaveCondOnImage): | |
| def __init__(self, dataset_cfg): | |
| train_transforms = transforms.Compose([ | |
| Resize3D(128), | |
| RandomResizedCrop3D(112, scale=(0.5, 1.0)), | |
| RandomHorizontalFlip3D(), | |
| ColorJitter3D(brightness=0.4, saturation=0.4, contrast=0.2, hue=0.1), | |
| ToTensor3D(), | |
| Normalize3D(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]), | |
| ]) | |
| super().__init__('train', frame_transforms=train_transforms, **dataset_cfg) | |
| class GreatestHitWaveCondOnImageValidation(GreatestHitWaveCondOnImage): | |
| def __init__(self, dataset_cfg): | |
| valid_transforms = transforms.Compose([ | |
| Resize3D(128), | |
| CenterCrop3D(112), | |
| ToTensor3D(), | |
| Normalize3D(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]), | |
| ]) | |
| super().__init__('val', frame_transforms=valid_transforms, **dataset_cfg) | |
| class GreatestHitWaveCondOnImageTest(GreatestHitWaveCondOnImage): | |
| def __init__(self, dataset_cfg): | |
| test_transforms = transforms.Compose([ | |
| Resize3D(128), | |
| CenterCrop3D(112), | |
| ToTensor3D(), | |
| Normalize3D(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]), | |
| ]) | |
| super().__init__('test', frame_transforms=test_transforms, **dataset_cfg) | |
| def draw_spec(spec, dest, cmap='magma'): | |
| plt.imshow(spec, cmap=cmap, origin='lower') | |
| plt.axis('off') | |
| plt.savefig(dest, bbox_inches='tight', pad_inches=0., dpi=300) | |
| plt.close() | |
| if __name__ == '__main__': | |
| import sys | |
| from omegaconf import OmegaConf | |
| # cfg = OmegaConf.load('configs/greatesthit_transformer_with_vNet_randshift_2s_GH_vqgan_no_earlystop.yaml') | |
| cfg = OmegaConf.load('configs/greatesthit_codebook.yaml') | |
| data = instantiate_from_config(cfg.data) | |
| data.prepare_data() | |
| data.setup() | |
| print(len(data.datasets['train'])) | |
| print(data.datasets['train'][24]) | |