File size: 1,876 Bytes
df07554 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | from scipy.misc import imsave
import dlib
import os
import glob
import numpy as np
import cv2
from multiprocessing import Pool
import pdb
from torch.utils.data import DataLoader, Dataset
import time
class MyDataset(Dataset):
def __init__(self):
self.IN = 'GRID/'
self.OUT = 'GRID_imgs/'
self.wav = 'GRID_wavs/'
with open('GRID_files.txt', 'r') as f:
files = [line.strip() for line in f.readlines()]
self.files = []
for file in files:
_, ext = os.path.splitext(file)
if ext == '.XML': continue
self.files.append(file)
print(file)
wav = file.replace(self.IN, self.wav).replace(ext, '.wav')
path = os.path.split(wav)[0]
if not os.path.exists(path):
os.makedirs(path)
def __len__(self):
return len(self.files)
def __getitem__(self, idx):
file = self.files[idx]
_, ext = os.path.splitext(file)
dst = file.replace(self.IN, self.OUT).replace(ext, '')
if not os.path.exists(dst):
os.makedirs(dst)
cmd = 'ffmpeg -i \'{}\' -qscale:v 2 -r 25 \'{}/%d.jpg\''.format(file, dst)
os.system(cmd)
wav = file.replace(self.IN, self.wav).replace(ext, '.wav')
cmd = 'ffmpeg -y -i \'{}\' -async 1 -ac 1 -vn -acodec pcm_s16le -ar 16000 \'{}\' '.format(file, wav)
os.system(cmd)
return dst
if __name__ == '__main__':
dataset = MyDataset()
loader = DataLoader(
dataset, num_workers=32, batch_size=128,
shuffle=False, drop_last=False
)
tic = time.time()
for (i, batch) in enumerate(loader):
eta = (1.0*time.time()-tic)/(i+1) * (len(loader)-i)
print('eta:{}'.format(eta/3600.0))
|