foleycrafter/models/specvqgan/data/impactset.py

import json
import os
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
from glob import glob
import sys
import soundfile
import copy
import csv
import noisereduce as nr

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

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()

def convert_to_decibel(arr):
    ref = 1
    return 20 * np.log10(abs(arr + 1e-4) / ref)

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 ImpactSetWave(torch.utils.data.Dataset):

    def __init__(self, split, random_crop, mel_num, spec_crop_len,
                L=2.0, denoise=False, splits_path='./data',
                data_path='data/ImpactSet/impactset-proccess-resize'):
        super().__init__()
        self.split = split
        self.splits_path = splits_path
        self.data_path = data_path
        self.L = L
        self.denoise = denoise

        video_name_split_path = os.path.join(splits_path, f'countixAV_{split}.json')
        if not os.path.exists(video_name_split_path):
            self.make_split_files()
        video_name = json.load(open(video_name_split_path, 'r'))
        self.video_frame_cnt = {v: len(os.listdir(os.path.join(self.data_path, v, 'frames'))) 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}_resampled.wav') for v in video_name}
        self.dataset = video_name

        self.wav_transforms = transforms.Compose([
            MakeMono(),
            Padding(target_len=int(SR * self.L)),
        ])
        
        self.spec_transforms = CropImage([mel_num, spec_crop_len], random_crop)

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

    def __getitem__(self, idx):
        item = {}
        video = self.dataset[idx]

        available_frame_idx = self.video_frame_cnt[video] - self.left_over
        wav = None
        spec = None
        max_db = -np.inf
        wave_path = ''
        cur_wave_path = self.video_audio_path[video]
        if self.denoise:
            cur_wave_path = cur_wave_path.replace('.wav', '_denoised.wav')
        for _ in range(10):
            start_idx = torch.randint(0, available_frame_idx, (1,)).tolist()[0]
            # target
            start_t = (start_idx + 0.5) / FPS
            start_audio_idx = non_negative(start_t * SR)

            cur_wav, _ = soundfile.read(cur_wave_path, frames=int(SR * self.L), start=start_audio_idx)

            decibel = convert_to_decibel(cur_wav)
            if float(np.mean(decibel)) > max_db:
                wav = cur_wav
                wave_path = cur_wave_path
                max_db = float(np.mean(decibel))
            if max_db >= -40:
                break

        # print(max_db)
        wav = self.wav_transforms(wav)
        item['image'] = wav # (80, 173)
        # item['wav'] = wav
        item['file_path_wav_'] = wave_path

        item['label'] = 'None'
        item['target'] = 'None'

        return item

    def make_split_files(self):
        raise NotImplementedError

class ImpactSetWaveTrain(ImpactSetWave):
    def __init__(self, specs_dataset_cfg):
        super().__init__('train', **specs_dataset_cfg)

class ImpactSetWaveValidation(ImpactSetWave):
    def __init__(self, specs_dataset_cfg):
        super().__init__('val', **specs_dataset_cfg)

class ImpactSetWaveTest(ImpactSetWave):
    def __init__(self, specs_dataset_cfg):
        super().__init__('test', **specs_dataset_cfg)


class ImpactSetSpec(torch.utils.data.Dataset):

    def __init__(self, split, random_crop, mel_num, spec_crop_len,
                L=2.0, denoise=False, splits_path='./data',
                data_path='data/ImpactSet/impactset-proccess-resize'):
        super().__init__()
        self.split = split
        self.splits_path = splits_path
        self.data_path = data_path
        self.L = L
        self.denoise = denoise

        video_name_split_path = os.path.join(splits_path, f'countixAV_{split}.json')
        if not os.path.exists(video_name_split_path):
            self.make_split_files()
        video_name = json.load(open(video_name_split_path, 'r'))
        self.video_frame_cnt = {v: len(os.listdir(os.path.join(self.data_path, v, 'frames'))) 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}_resampled.wav') for v in video_name}
        self.dataset = video_name

        self.wav_transforms = transforms.Compose([
            MakeMono(),
            SpectrogramTorchAudio(nfft=1024, hoplen=1024//4, spec_power=1),
            MelScaleTorchAudio(sr=SR, stft=513, fmin=125, fmax=7600, nmels=80),
            LowerThresh(1e-5),
            Log10(),
            Multiply(20),
            Subtract(20),
            Add(100),
            Divide(100),
            Clip(0, 1.0),
            TrimSpec(173),
        ])
        
        self.spec_transforms = CropImage([mel_num, spec_crop_len], random_crop)

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

    def __getitem__(self, idx):
        item = {}
        video = self.dataset[idx]

        available_frame_idx = self.video_frame_cnt[video] - self.left_over
        wav = None
        spec = None
        max_rms = -np.inf
        wave_path = ''
        cur_wave_path = self.video_audio_path[video]
        if self.denoise:
            cur_wave_path = cur_wave_path.replace('.wav', '_denoised.wav')
        for _ in range(10):
            start_idx = torch.randint(0, available_frame_idx, (1,)).tolist()[0]
            # target
            start_t = (start_idx + 0.5) / FPS
            start_audio_idx = non_negative(start_t * SR)

            cur_wav, _ = soundfile.read(cur_wave_path, frames=int(SR * self.L), start=start_audio_idx)

            if self.wav_transforms is not None:
                spec_tensor = self.wav_transforms(torch.tensor(cur_wav).float())
                cur_spec = spec_tensor.numpy()
            # zeros padding if not enough spec t steps
            if cur_spec.shape[1] < 173:
                pad = np.zeros((80, 173), dtype=cur_spec.dtype)
                pad[:, :cur_spec.shape[1]] = cur_spec
                cur_spec = pad
            rms_val = rms(cur_spec)
            if rms_val > max_rms:
                wav = cur_wav
                spec = cur_spec
                wave_path = cur_wave_path
                max_rms = rms_val
            # print(rms_val)
            if max_rms >= 0.1:
                break

        item['image'] = 2 * spec - 1 # (80, 173)
        # item['wav'] = wav
        item['file_path_wav_'] = wave_path

        item['label'] = 'None'
        item['target'] = 'None'

        if self.spec_transforms is not None:
            item = self.spec_transforms(item)
        return item

    def make_split_files(self):
        raise NotImplementedError

class ImpactSetSpecTrain(ImpactSetSpec):
    def __init__(self, specs_dataset_cfg):
        super().__init__('train', **specs_dataset_cfg)

class ImpactSetSpecValidation(ImpactSetSpec):
    def __init__(self, specs_dataset_cfg):
        super().__init__('val', **specs_dataset_cfg)

class ImpactSetSpecTest(ImpactSetSpec):
    def __init__(self, specs_dataset_cfg):
        super().__init__('test', **specs_dataset_cfg)



class ImpactSetWaveTestTime(torch.utils.data.Dataset):

    def __init__(self, split, random_crop, mel_num, spec_crop_len,
                L=2.0, denoise=False, splits_path='./data',
                data_path='data/ImpactSet/impactset-proccess-resize'):
        super().__init__()
        self.split = split
        self.splits_path = splits_path
        self.data_path = data_path
        self.L = L
        self.denoise = denoise

        self.video_list = glob('data/ImpactSet/RawVideos/StockVideo_sound/*.wav') + [
            'data/ImpactSet/RawVideos/YouTube-impact-ccl/1_ckbCU5aQs/1_ckbCU5aQs_0013_0016_resize.wav', 
            'data/ImpactSet/RawVideos/YouTube-impact-ccl/GFmuVBiwz6k/GFmuVBiwz6k_0034_0054_resize.wav',
            'data/ImpactSet/RawVideos/YouTube-impact-ccl/OsPcY316h1M/OsPcY316h1M_0000_0005_resize.wav',
            'data/ImpactSet/RawVideos/YouTube-impact-ccl/SExIpBIBj_k/SExIpBIBj_k_0009_0019_resize.wav',
            'data/ImpactSet/RawVideos/YouTube-impact-ccl/S6TkbV4B4QI/S6TkbV4B4QI_0028_0036_resize.wav',
            'data/ImpactSet/RawVideos/YouTube-impact-ccl/2Ld24pPIn3k/2Ld24pPIn3k_0005_0011_resize.wav',
            'data/ImpactSet/RawVideos/YouTube-impact-ccl/6d1YS7fdBK4/6d1YS7fdBK4_0007_0019_resize.wav',
            'data/ImpactSet/RawVideos/YouTube-impact-ccl/JnBsmJgEkiw/JnBsmJgEkiw_0008_0016_resize.wav',
            'data/ImpactSet/RawVideos/YouTube-impact-ccl/xcUyiXt0gjo/xcUyiXt0gjo_0015_0021_resize.wav',
            'data/ImpactSet/RawVideos/YouTube-impact-ccl/4DRFJnZjpMM/4DRFJnZjpMM_0000_0010_resize.wav'
        ] + glob('data/ImpactSet/RawVideos/self_recorded/*_resize.wav')

        self.wav_transforms = transforms.Compose([
            MakeMono(),
            SpectrogramTorchAudio(nfft=1024, hoplen=1024//4, spec_power=1),
            MelScaleTorchAudio(sr=SR, stft=513, fmin=125, fmax=7600, nmels=80),
            LowerThresh(1e-5),
            Log10(),
            Multiply(20),
            Subtract(20),
            Add(100),
            Divide(100),
            Clip(0, 1.0),
            TrimSpec(173),
        ])
        self.spec_transforms = CropImage([mel_num, spec_crop_len], random_crop)

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

    def __getitem__(self, idx):
        item = {}

        wave_path = self.video_list[idx]

        wav, _ = soundfile.read(wave_path)
        start_idx = random.randint(0, min(4, wav.shape[0] - int(SR * self.L)))
        wav = wav[start_idx:start_idx+int(SR * self.L)]

        if self.denoise:
            if len(wav.shape) == 1:
                wav = wav[None, :]
            wav = nr.reduce_noise(y=wav, sr=SR, n_fft=1024, hop_length=1024//4)
            wav = wav.squeeze()
        if self.wav_transforms is not None:
            spec_tensor = self.wav_transforms(torch.tensor(wav).float())
            spec = spec_tensor.numpy()
        if spec.shape[1] < 173:
            pad = np.zeros((80, 173), dtype=spec.dtype)
            pad[:, :spec.shape[1]] = spec
            spec = pad

        item['image'] = 2 * spec - 1 # (80, 173)
        # item['wav'] = wav
        item['file_path_wav_'] = wave_path

        item['label'] = 'None'
        item['target'] = 'None'

        if self.spec_transforms is not None:
            item = self.spec_transforms(item)
        return item

    def make_split_files(self):
        raise NotImplementedError

class ImpactSetWaveTestTimeTrain(ImpactSetWaveTestTime):
    def __init__(self, specs_dataset_cfg):
        super().__init__('train', **specs_dataset_cfg)

class ImpactSetWaveTestTimeValidation(ImpactSetWaveTestTime):
    def __init__(self, specs_dataset_cfg):
        super().__init__('val', **specs_dataset_cfg)

class ImpactSetWaveTestTimeTest(ImpactSetWaveTestTime):
    def __init__(self, specs_dataset_cfg):
        super().__init__('test', **specs_dataset_cfg)


class ImpactSetWaveWithSilent(torch.utils.data.Dataset):

    def __init__(self, split, random_crop, mel_num, spec_crop_len,
                L=2.0, denoise=False, splits_path='./data',
                data_path='data/ImpactSet/impactset-proccess-resize'):
        super().__init__()
        self.split = split
        self.splits_path = splits_path
        self.data_path = data_path
        self.L = L
        self.denoise = denoise

        video_name_split_path = os.path.join(splits_path, f'countixAV_{split}.json')
        if not os.path.exists(video_name_split_path):
            self.make_split_files()
        video_name = json.load(open(video_name_split_path, 'r'))
        self.video_frame_cnt = {v: len(os.listdir(os.path.join(self.data_path, v, 'frames'))) 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}_resampled.wav') for v in video_name}
        self.dataset = video_name

        self.wav_transforms = transforms.Compose([
            MakeMono(),
            Padding(target_len=int(SR * self.L)),
        ])
        
        self.spec_transforms = CropImage([mel_num, spec_crop_len], random_crop)

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

    def __getitem__(self, idx):
        item = {}
        video = self.dataset[idx]

        available_frame_idx = self.video_frame_cnt[video] - self.left_over
        wave_path = self.video_audio_path[video]
        if self.denoise:
            wave_path = wave_path.replace('.wav', '_denoised.wav')
        start_idx = torch.randint(0, available_frame_idx, (1,)).tolist()[0]
        # target
        start_t = (start_idx + 0.5) / FPS
        start_audio_idx = non_negative(start_t * SR)

        wav, _ = soundfile.read(wave_path, frames=int(SR * self.L), start=start_audio_idx)

        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

    def make_split_files(self):
        raise NotImplementedError

class ImpactSetWaveWithSilentTrain(ImpactSetWaveWithSilent):
    def __init__(self, specs_dataset_cfg):
        super().__init__('train', **specs_dataset_cfg)

class ImpactSetWaveWithSilentValidation(ImpactSetWaveWithSilent):
    def __init__(self, specs_dataset_cfg):
        super().__init__('val', **specs_dataset_cfg)

class ImpactSetWaveWithSilentTest(ImpactSetWaveWithSilent):
    def __init__(self, specs_dataset_cfg):
        super().__init__('test', **specs_dataset_cfg)


class ImpactSetWaveCondOnImage(torch.utils.data.Dataset):

    def __init__(self, split,
                L=2.0, frame_transforms=None, denoise=False, splits_path='./data',
                data_path='data/ImpactSet/impactset-proccess-resize',
                p_outside_cond=0.):
        super().__init__()
        self.split = split
        self.splits_path = splits_path
        self.frame_transforms = frame_transforms
        self.data_path = data_path
        self.L = L
        self.denoise = denoise
        self.p_outside_cond = torch.tensor(p_outside_cond)

        video_name_split_path = os.path.join(splits_path, f'countixAV_{split}.json')
        if not os.path.exists(video_name_split_path):
            self.make_split_files()
        video_name = json.load(open(video_name_split_path, 'r'))
        self.video_frame_cnt = {v: len(os.listdir(os.path.join(self.data_path, v, 'frames'))) for v in video_name}
        self.left_over = int(FPS * L + 1)
        for v, cnt in self.video_frame_cnt.items():
            if cnt - (3*self.left_over) <= 0:
                video_name.remove(v)
        self.video_audio_path = {v: os.path.join(self.data_path, v, f'audio/{v}_resampled.wav') for v in video_name}
        self.dataset = video_name

        video_timing_split_path = os.path.join(splits_path, f'countixAV_{split}_timing.json')
        self.video_timing = json.load(open(video_timing_split_path, 'r'))
        self.video_timing = {v: [int(float(t) * FPS) for t in ts] for v, ts in self.video_timing.items()}

        if split != 'test':
            video_class_path = os.path.join(splits_path, f'countixAV_{split}_class.json')
            if not os.path.exists(video_class_path):
                self.make_video_class()
            self.video_class = json.load(open(video_class_path, 'r'))
            self.class2video = {}
            for v, c in self.video_class.items():
                if c not in self.class2video.keys():
                    self.class2video[c] = []
                self.class2video[c].append(v)

        self.wav_transforms = transforms.Compose([
            MakeMono(),
            Padding(target_len=int(SR * self.L)),
        ])
        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]),
            ])

    def make_video_class(self):
        meta_path = f'data/ImpactSet/data-info/CountixAV_{self.split}.csv'
        video_class = {}
        with open(meta_path, 'r') as f:
            reader = csv.reader(f)
            for i, row in enumerate(reader):
                if i == 0:
                    continue
                vid, k_st, k_et = row[:3]
                video_name = f'{vid}_{int(k_st):0>4d}_{int(k_et):0>4d}'
                if video_name not in self.dataset:
                    continue
                video_class[video_name] = row[-1]
        with open(os.path.join(self.splits_path, f'countixAV_{self.split}_class.json'), 'w') as f:
            json.dump(video_class, f)

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

    def __getitem__(self, idx):
        item = {}
        video = self.dataset[idx]

        available_frame_idx = self.video_frame_cnt[video] - self.left_over
        rep_start_idx, rep_end_idx = self.video_timing[video]
        rep_end_idx = min(available_frame_idx, rep_end_idx)
        if available_frame_idx <= rep_start_idx + self.L * FPS:
            idx_set = list(range(0, available_frame_idx))
        else:
            idx_set = list(range(rep_start_idx, rep_end_idx))
        start_idx = sample(idx_set, k=1)[0]

        wave_path = self.video_audio_path[video]
        if self.denoise:
            wave_path = wave_path.replace('.wav', '_denoised.wav')

        # target
        start_t = (start_idx + 0.5) / FPS
        end_idx= non_negative(start_idx + FPS * self.L)
        start_audio_idx = non_negative(start_t * SR)
        wav, sr = soundfile.read(wave_path, frames=int(SR * self.L), start=start_audio_idx)
        assert sr == SR
        wav = self.wav_transforms(wav)
        frame_path = os.path.join(self.data_path, video, 'frames')
        frames = [Image.open(os.path.join(
            frame_path, f'frame{i+1:0>6d}.jpg')).convert('RGB') for i in
            range(start_idx, end_idx)]

        if torch.all(torch.bernoulli(self.p_outside_cond) == 1.) and self.split != 'test':
            # outside from the same class
            cur_class = self.video_class[video]
            tmp_video = copy.copy(self.class2video[cur_class])
            if len(tmp_video) > 1:
                # if only 1 video in the class, use itself
                tmp_video.remove(video)
            cond_video = sample(tmp_video, k=1)[0]
            cond_available_frame_idx = self.video_frame_cnt[cond_video] - self.left_over
            cond_start_idx = torch.randint(0, cond_available_frame_idx, (1,)).tolist()[0]
        else:
            cond_video = video
            idx_set = list(range(0, start_idx)) + list(range(end_idx, available_frame_idx))
            cond_start_idx = random.sample(idx_set, k=1)[0]

        cond_end_idx = non_negative(cond_start_idx + FPS * self.L)
        cond_start_t = (cond_start_idx + 0.5) / FPS
        cond_audio_idx = non_negative(cond_start_t * SR)
        cond_frame_path = os.path.join(self.data_path, cond_video, 'frames')
        cond_wave_path = self.video_audio_path[cond_video]

        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_idx, cond_end_idx)]
        cond_wav, sr = soundfile.read(cond_wave_path, frames=int(SR * self.L), start=cond_audio_idx)
        assert sr == SR
        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 make_split_files(self):
        raise NotImplementedError


class ImpactSetWaveCondOnImageTrain(ImpactSetWaveCondOnImage):
    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 ImpactSetWaveCondOnImageValidation(ImpactSetWaveCondOnImage):
    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 ImpactSetWaveCondOnImageTest(ImpactSetWaveCondOnImage):
    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 ImpactSetCleanWaveCondOnImage(ImpactSetWaveCondOnImage):
    def __init__(self, split, L=2, frame_transforms=None, denoise=False, splits_path='./data', data_path='data/ImpactSet/impactset-proccess-resize', p_outside_cond=0):
        super().__init__(split, L, frame_transforms, denoise, splits_path, data_path, p_outside_cond)
        pred_timing_path = f'data/countixAV_{split}_timing_processed_0.20.json'
        assert os.path.exists(pred_timing_path)
        self.pred_timing = json.load(open(pred_timing_path, 'r'))

        self.dataset = []
        for v, ts in self.pred_timing.items():
            if v in self.video_audio_path.keys():
                for t in ts:
                    self.dataset.append([v, t])

    def __getitem__(self, idx):
        item = {}
        video, start_t = self.dataset[idx]
        available_frame_idx = self.video_frame_cnt[video] - self.left_over
        available_timing = (available_frame_idx + 0.5) / FPS
        start_t = float(start_t)
        start_t = min(start_t, available_timing)

        start_idx = non_negative(start_t * FPS - 0.5)

        wave_path = self.video_audio_path[video]
        if self.denoise:
            wave_path = wave_path.replace('.wav', '_denoised.wav')

        # target
        end_idx= non_negative(start_idx + FPS * self.L)
        start_audio_idx = non_negative(start_t * SR)
        wav, sr = soundfile.read(wave_path, frames=int(SR * self.L), start=start_audio_idx)
        assert sr == SR
        wav = self.wav_transforms(wav)
        frame_path = os.path.join(self.data_path, video, 'frames')
        frames = [Image.open(os.path.join(
            frame_path, f'frame{i+1:0>6d}.jpg')).convert('RGB') for i in
            range(start_idx, end_idx)]

        if torch.all(torch.bernoulli(self.p_outside_cond) == 1.):
            other_video = list(self.pred_timing.keys())
            other_video.remove(video)
            cond_video = sample(other_video, k=1)[0]
            cond_available_frame_idx = self.video_frame_cnt[cond_video] - self.left_over
            cond_available_timing = (cond_available_frame_idx + 0.5) / FPS
        else:
            cond_video = video
            cond_available_timing = available_timing

        cond_start_t = sample(self.pred_timing[cond_video], k=1)[0]
        cond_start_t = float(cond_start_t)
        cond_start_t = min(cond_start_t, cond_available_timing)
        cond_start_idx = non_negative(cond_start_t * FPS - 0.5)
        cond_end_idx = non_negative(cond_start_idx + FPS * self.L)
        cond_audio_idx = non_negative(cond_start_t * SR)
        cond_frame_path = os.path.join(self.data_path, cond_video, 'frames')
        cond_wave_path = self.video_audio_path[cond_video]

        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_idx, cond_end_idx)]
        cond_wav, sr = soundfile.read(cond_wave_path, frames=int(SR * self.L), start=cond_audio_idx)
        assert sr == SR
        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


class ImpactSetCleanWaveCondOnImageTrain(ImpactSetCleanWaveCondOnImage):
    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 ImpactSetCleanWaveCondOnImageValidation(ImpactSetCleanWaveCondOnImage):
    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 ImpactSetCleanWaveCondOnImageTest(ImpactSetCleanWaveCondOnImage):
    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)


if __name__ == '__main__':
    import sys

    from omegaconf import OmegaConf
    cfg = OmegaConf.load('configs/countixAV_transformer_denoise_clean.yaml')
    data = instantiate_from_config(cfg.data)
    data.prepare_data()
    data.setup()

    print(data.datasets['train'])
    print(len(data.datasets['train']))
    # print(data.datasets['train'][24])
    exit()

    stats = []
    torch.manual_seed(0)
    np.random.seed(0)
    random.seed = 0
    for k in range(1):
        x = np.arange(SR * 2)
        for i in tqdm(range(len(data.datasets['train']))):
            wav = data.datasets['train'][i]['wav']
            spec = data.datasets['train'][i]['image']
            spec = 0.5 * (spec + 1)
            spec_rms = rms(spec)
            stats.append(float(spec_rms))
            # plt.plot(x, wav)
            # plt.ylim(-1, 1)
            # plt.savefig(f'tmp/th0.1_wav_e_{k}_{i}_{mean_val:.3f}_{spec_rms:.3f}.png')
            # plt.close()
            # plt.cla()
            soundfile.write(f'tmp/wav_e_{k}_{i}_{spec_rms:.3f}.wav', wav, SR)
            draw_spec(spec, f'tmp/wav_spec_e_{k}_{i}_{spec_rms:.3f}.png')
            if i == 100:
                break
    # plt.hist(stats, bins=50)
    # plt.savefig(f'tmp/rms_spec_stats.png')