audio_diffusion_attacks_forhf/tests/test_datasets.py

""" 
test_datasets.py
    Desc: a file for testing datasets
"""

import numpy as np 
import pandas as pd 
import matplotlib.pyplot as plt 
import seaborn as sns 
import glob 

import sys 
sys.path.append('./')

from models.datasets.music_datasets import MusicMelDataset

def test_MusicMelDataset():
    audio_files = glob.glob("/data/robbizorg/music_datasets/fma/data/fma_processed/*/*.npy")

    dataset = MusicMelDataset(audio_files)

    failed_files = []
    # See if errors happen
    for i in range(len(dataset)):
        try:
            dataset.__getitem__(i)
        except Exception as e:
            # print(i, repr(e))
            failed_files.append(dataset.audio_files[i])
    
    print("Len of Failed Files: ", len(failed_files))

    final_files = set(audio_files).difference(set(failed_files))

    with open("./data/fma_valid_files.npy", 'wb') as f:
        np.save(f, final_files)

if __name__ == '__main__':
    test_MusicMelDataset()