Spaces:
Running
on
Zero
Running
on
Zero
File size: 6,562 Bytes
a84a65c |
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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 |
import collections
import csv
import logging
import os
import random
from glob import glob
from pathlib import Path
import numpy as np
import torch
import torchvision
logger = logging.getLogger(f'main.{__name__}')
class VGGSound(torch.utils.data.Dataset):
def __init__(self, split, specs_dir, transforms=None, splits_path='./data', meta_path='./data/vggsound.csv'):
super().__init__()
self.split = split
self.specs_dir = specs_dir
self.transforms = transforms
self.splits_path = splits_path
self.meta_path = meta_path
vggsound_meta = list(csv.reader(open(meta_path), quotechar='"'))
unique_classes = sorted(list(set(row[2] for row in vggsound_meta)))
self.label2target = {label: target for target, label in enumerate(unique_classes)}
self.target2label = {target: label for label, target in self.label2target.items()}
self.video2target = {row[0]: self.label2target[row[2]] for row in vggsound_meta}
split_clip_ids_path = os.path.join(splits_path, f'vggsound_{split}.txt')
if not os.path.exists(split_clip_ids_path):
self.make_split_files()
clip_ids_with_timestamp = open(split_clip_ids_path).read().splitlines()
clip_paths = [os.path.join(specs_dir, v + '_mel.npy') for v in clip_ids_with_timestamp]
self.dataset = clip_paths
# self.dataset = clip_paths[:10000] # overfit one batch
# 'zyTX_1BXKDE_16000_26000'[:11] -> 'zyTX_1BXKDE'
vid_classes = [self.video2target[Path(path).stem[:11]] for path in self.dataset]
class2count = collections.Counter(vid_classes)
self.class_counts = torch.tensor([class2count[cls] for cls in range(len(class2count))])
# self.sample_weights = [len(self.dataset) / class2count[self.video2target[Path(path).stem[:11]]] for path in self.dataset]
def __getitem__(self, idx):
item = {}
spec_path = self.dataset[idx]
# 'zyTX_1BXKDE_16000_26000' -> 'zyTX_1BXKDE'
video_name = Path(spec_path).stem[:11]
item['input'] = np.load(spec_path)
item['input_path'] = spec_path
# if self.split in ['train', 'valid']:
item['target'] = self.video2target[video_name]
item['label'] = self.target2label[item['target']]
if self.transforms is not None:
item = self.transforms(item)
return item
def __len__(self):
return len(self.dataset)
def make_split_files(self):
random.seed(1337)
logger.info(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_vid_paths = sorted(glob(os.path.join(self.specs_dir, '*_mel.npy')))
logger.info(f'The number of clips available after download: {len(available_vid_paths)}')
# original (full) train and test sets
vggsound_meta = list(csv.reader(open(self.meta_path), quotechar='"'))
train_vids = {row[0] for row in vggsound_meta if row[3] == 'train'}
test_vids = {row[0] for row in vggsound_meta if row[3] == 'test'}
logger.info(f'The number of videos in vggsound train set: {len(train_vids)}')
logger.info(f'The number of videos in vggsound test set: {len(test_vids)}')
# class counts in test set. We would like to have the same distribution in valid
unique_classes = sorted(list(set(row[2] for row in vggsound_meta)))
label2target = {label: target for target, label in enumerate(unique_classes)}
video2target = {row[0]: label2target[row[2]] for row in vggsound_meta}
test_vid_classes = [video2target[vid] for vid in test_vids]
test_target2count = collections.Counter(test_vid_classes)
# now given the counts from test set, sample the same count for validation and the rest leave in train
train_vids_wo_valid, valid_vids = set(), set()
for target, label in enumerate(label2target.keys()):
class_train_vids = [vid for vid in train_vids if video2target[vid] == target]
random.shuffle(class_train_vids)
count = test_target2count[target]
valid_vids.update(class_train_vids[:count])
train_vids_wo_valid.update(class_train_vids[count:])
# make file with a list of available test videos (each video should contain timestamps as well)
train_i = valid_i = test_i = 0
with open(os.path.join(self.splits_path, 'vggsound_train.txt'), 'w') as train_file, \
open(os.path.join(self.splits_path, 'vggsound_valid.txt'), 'w') as valid_file, \
open(os.path.join(self.splits_path, 'vggsound_test.txt'), 'w') as test_file:
for path in available_vid_paths:
path = path.replace('_mel.npy', '')
vid_name = Path(path).name
# 'zyTX_1BXKDE_16000_26000'[:11] -> 'zyTX_1BXKDE'
if vid_name[:11] in train_vids_wo_valid:
train_file.write(vid_name + '\n')
train_i += 1
elif vid_name[:11] in valid_vids:
valid_file.write(vid_name + '\n')
valid_i += 1
elif vid_name[:11] in test_vids:
test_file.write(vid_name + '\n')
test_i += 1
else:
raise Exception(f'Clip {vid_name} is neither in train, valid nor test. Strange.')
logger.info(f'Put {train_i} clips to the train set and saved it to ./data/vggsound_train.txt')
logger.info(f'Put {valid_i} clips to the valid set and saved it to ./data/vggsound_valid.txt')
logger.info(f'Put {test_i} clips to the test set and saved it to ./data/vggsound_test.txt')
if __name__ == '__main__':
from transforms import Crop, StandardNormalizeAudio, ToTensor
specs_path = '/home/nvme/data/vggsound/features/melspec_10s_22050hz/'
transforms = torchvision.transforms.transforms.Compose([
StandardNormalizeAudio(specs_path),
ToTensor(),
Crop([80, 848]),
])
datasets = {
'train': VGGSound('train', specs_path, transforms),
'valid': VGGSound('valid', specs_path, transforms),
'test': VGGSound('test', specs_path, transforms),
}
print(datasets['train'][0])
print(datasets['valid'][0])
print(datasets['test'][0])
print(datasets['train'].class_counts)
print(datasets['valid'].class_counts)
print(datasets['test'].class_counts)
|