SonicDiffusion / CLAP /msclap /esc50_dataset.py
root
init
9778d56
raw
history blame
No virus
2.79 kB
from torch.utils.data import Dataset
from torchvision.datasets.utils import download_url
from tqdm import tqdm
import pandas as pd
import os
import torch.nn as nn
import torch
class AudioDataset(Dataset):
def __init__(self, root: str, download: bool = True):
self.root = os.path.expanduser(root)
if download:
self.download()
def __getitem__(self, index):
raise NotImplementedError
def download(self):
raise NotImplementedError
def __len__(self):
raise NotImplementedError
class ESC50(AudioDataset):
base_folder = 'ESC-50-master'
url = "https://github.com/karolpiczak/ESC-50/archive/refs/heads/master.zip"
filename = "ESC-50-master.zip"
num_files_in_dir = 2000
audio_dir = 'audio'
label_col = 'category'
file_col = 'filename'
meta = {
'filename': os.path.join('meta','esc50.csv'),
}
def __init__(self, root, reading_transformations: nn.Module = None, download: bool = True):
super().__init__(root)
self._load_meta()
self.targets, self.audio_paths = [], []
self.pre_transformations = reading_transformations
print("Loading audio files")
# self.df['filename'] = os.path.join(self.root, self.base_folder, self.audio_dir) + os.sep + self.df['filename']
self.df['category'] = self.df['category'].str.replace('_',' ')
for _, row in tqdm(self.df.iterrows()):
file_path = os.path.join(self.root, self.base_folder, self.audio_dir, row[self.file_col])
self.targets.append(row[self.label_col])
self.audio_paths.append(file_path)
def _load_meta(self):
path = os.path.join(self.root, self.base_folder, self.meta['filename'])
self.df = pd.read_csv(path)
self.class_to_idx = {}
self.classes = [x.replace('_',' ') for x in sorted(self.df[self.label_col].unique())]
for i, category in enumerate(self.classes):
self.class_to_idx[category] = i
def __getitem__(self, index):
"""
Args:
index (int): Index
Returns:
tuple: (image, target) where target is index of the target class.
"""
file_path, target = self.audio_paths[index], self.targets[index]
idx = torch.tensor(self.class_to_idx[target])
one_hot_target = torch.zeros(len(self.classes)).scatter_(0, idx, 1).reshape(1,-1)
return file_path, target, one_hot_target
def __len__(self):
return len(self.audio_paths)
def download(self):
download_url(self.url, self.root, self.filename)
# extract file
from zipfile import ZipFile
with ZipFile(os.path.join(self.root, self.filename), 'r') as zip:
zip.extractall(path=self.root)