import torch
import torchaudio
import torchvision
import os
from tqdm import tqdm


p = '/home/kinetics/ActivityNet/Crawler/Kinetics/test_data/archery/002VmnaNvh4_000003_000013.mp4'
sr = 44100
audio_transform = torchaudio.transforms.MelSpectrogram(sample_rate=sr, n_mels=40)
phase = 'train' # train, valid, test


avgpool = torch.nn.AvgPool2d((2,2))
datas = []
batchcount=0
for name in tqdm(os.listdir('%s' % phase)):
    if name[-3:]=='.pt':
        # print('at '+name)
        f = torch.load('%s/%s' % (phase, name))
        for tensors in f:
            a = tensors[1] # 1, 152576, values btw -1 and 1
            if len(a) != 0:
                if len(a) == 1:
                    a = a[0]
                elif len(a) == 2:
                    a = a.mean(0)
                a -= a.min()
                if a.min() != a.max():
                    a /= a.max()
                    a = 2*a - 1 # values in [-1, 1]
                spec = audio_transform(a)
                spec = spec + 1e-10
                spec = spec.log()

                v = tensors[0]
                l = tensors[2]
                v = ((v.float()/255.0)-torch.FloatTensor([0.43216, 0.394666, 0.37645]))/torch.FloatTensor([0.22803, 0.22145, 0.216989])
                v = v.transpose(0,3).transpose(3,1)
                v = avgpool(v)
                
                datas.append((v,spec,l))

                if len(datas)==100:
                    torch.save(datas,'%s/batch_37%s.pdt' % (phase, batchcount))
                    batchcount+=1
                    datas=[]

torch.save(datas,'%s/batch_37%s.pdt' % (phase, batchcount))

