standardization.py

from matplotlib import pyplot as plt
from glob import glob
import os
import argparse
import shutil
import tqdm
import random
from pathlib import Path
import numpy as np
from tqdm import tqdm
import librosa
import pandas as pd
import shutil
from distutils.dir_util import copy_tree

from resampy import resample
from scipy.io import wavfile
from pydub import AudioSegment
import soundfile as sf

from audio_utils import (trim_relative_silence_from_audio,
                        lufs_norm)

# Define whether each fsd50k classes fall into background (2), foreground (1), or not usable (0)
fsd50k_classes = [1,0,0,2,2,1,1,1,0,0,1,1,1,1,2,1,2,1,1,0,0,
                  0,1,2,1,1,2,2,1,2,2,1,1,0,0,1,1,0,0,1,1,1,
                  1,0,0,0,1,1,0,1,1,2,1,1,0,1,1,0,1,1,2,2,1,
                  1,1,1,1,0,0,0,2,1,1,1,0,0,1,1,2,1,2,2,1,2,
                  0,0,1,0,0,1,0,2,1,1,1,1,0,0,0,2,0,0,2,2,0,
                  1,1,0,1,2,0,0,0,0,2,2,1,2,2,2,0,0,2,0,1,0,
                  0,0,1,1,1,1,2,2,2,2,1,1,0,0,0,1,1,1,1,0,0,
                  1,0,0,0,2,2,1,1,1,0,0,0,0,1,1,2,0,2,0,0,1,
                  1,1,1,2,2,1,1,1,1,2,2,2,2,0,1,1,2,1,2,1,2,
                  2,0,1,2,2,0,0,1,1,0,1]

def prepare_fsd50k(src_dir, dst_dir):

    '''
    Write the FSD50K dataset from src to dst directory
    Format the output as instructed, which, in this case, solely involves silence trimming

    Args:
        src_dir: root path of FSD50K 
        dst_dir: destination of formatted FSD50K
    '''
    out_sr = 44100

    dst_dir = Path(dst_dir)
    dst_dir.mkdir(parents=True, exist_ok= True)

    # First, let's create the ground truth-combined .csv
    df_dev = pd.read_csv(os.path.join(src_dir,'FSD50K.ground_truth/dev.csv'))
    df_eval = pd.read_csv(os.path.join(src_dir,'FSD50K.ground_truth/eval.csv'))
    df_eval['split'] = 'test'
    pd.concat([df_dev,df_eval]).to_csv(os.path.join(dst_dir,'all_combined.csv'))

    paths = glob(src_dir+'/**/*.wav', recursive=True)
    random.shuffle(paths)

    for path in tqdm(paths):
        data, _ = librosa.load(path, sr=out_sr)
        trim_s, trim_e = trim_relative_silence_from_audio(data, out_sr)
        out_name = os.path.basename(path)
        out_dir = Path(os.path.join(dst_dir,path.split('/')[-2]))
        out_dir.mkdir(parents=True, exist_ok= True)
        if not os.path.isfile(out_name):
            try:     
                sf.write(os.path.join(out_dir,out_name), data[trim_s:trim_e], out_sr, 'PCM_16')
            except Exception as e: 
                print(e)

def refactor_fsd50k(ground_truth_path, root_path):
    '''
    Given a FSD50K vocabulary list (by default found in FSD50K.ground_truth)
    Restructure the dataset folders according to the classes defined by "fsd50k_classes"

    Args:
        ground_truth_path: Dataset ground truth folder directory (by default found in FSD50K.ground_truth)
        root_path: root_path of the location where FSD50K will be formatted
    '''

    vocab_csv  = pd.read_csv(os.path.join(ground_truth_path,'vocabulary.csv'),names=["label", "group"])
    vocab_csv['class'] = fsd50k_classes
    for s in ['dev','eval']:
        class_paths = {1:Path(os.path.join(root_path,'foreground')),
                       2:Path(os.path.join(root_path,'background'))}
        labels_csv = pd.read_csv(os.path.join(ground_truth_path,'{}.csv'.format(s)))
        paths = glob(os.path.join(root_path,'FSD50K.{}_audio').format(s)+'/*.wav')

        for p in tqdm(paths):
            fname = os.path.basename(p).split('.')[0]
            ls = labels_csv.loc[labels_csv['fname'] == int(fname)]
            ls = ls['labels'].values[0].split(',')
            cs = []
            for l in ls:
                # check file class value
                cs.append(vocab_csv.loc[vocab_csv['label']==l]['class'].values[0])

            # move the file to the corresponding dir based on class recurrency
            rec_c = np.argmax(np.bincount(cs))
            if rec_c==0:
                os.remove(p)
                continue
            dst_dir = class_paths[rec_c]
 
            split = labels_csv.loc[labels_csv['fname']==int(fname)]['split'].values[0] if s=='dev' else s
            split_path = Path(os.path.join(dst_dir, split))
            split_path.mkdir(parents=True, exist_ok= True)
            shutil.move(p, split_path)

        os.rmdir(os.path.join(root_path,'FSD50K.{}_audio'.format(s)))
        copy_tree(ground_truth_path, root_path)

def prepare_fma(src_dir, dst_dir):

    '''
    Write the FMA dataset from src to dst directory
    Format the output as instructed. The original files are .mp3, and stereo. 
    Conversion to .wav-mono is thus needed

    Args:
        src_dir: root path of original FMA 
        dst_dir: destination of formatted FMA
    '''

    def get_split_from_fname(fname, df):
        '''Retrieve split from a given filename and metadata dataframe'''
        # Keep spit name consistent across datasets
        split_format = {'training':'train','validation':'val','test':'eval'}
        f_idx = int(fname.split('.')[0].lstrip("0"))
        return split_format[df.loc[f_idx]['set']]

    dst_dir = Path(dst_dir)
    dst_dir.mkdir(parents=True, exist_ok= True)

    # First we create the annotation csv for FMA
    create_gt_fma(os.path.join(src_dir,'fma_metadata'), dst_dir)

    out_sr = 44100
    paths = glob(src_dir+'/**/*.mp3')

    # Load metadata csv and filter out columns
    df_tracks = pd.read_csv(os.path.join(src_dir,'fma_metadata','tracks.csv'))
    df_tracks = df_tracks.loc[df_tracks['set.1'].isin(['small', 'medium'])]
    df_ids    = df_tracks.loc[2:][['Unnamed: 0','set']]
    df_ids['Unnamed: 0'] = pd.to_numeric(df_ids['Unnamed: 0'])
    df_ids    = df_ids.set_index('Unnamed: 0')

    for path in tqdm(paths):
        fname = os.path.basename(path).split('.')[0]+'.wav'
        split = get_split_from_fname(fname, df_ids)

        sub = path.split('/')[-2]
        full_dst = Path(os.path.join(dst_dir,split,sub))
        full_dst.mkdir(parents=True, exist_ok= True)

        out_name = os.path.join(full_dst, fname)
        if not os.path.isfile(out_name):
            try:                                                     
                sound = AudioSegment.from_mp3(path)
                # Mono, 44100, WAV-PCM
                sound = sound.set_frame_rate(out_sr)
                sound = sound.set_channels(1)
                data  = np.array(sound.get_array_of_samples())
                sf.write(out_name, data, out_sr, 'PCM_16')
            except Exception as e: 
                print(e)

def create_gt_fma(fma_meta_folder, dst_dir):
    '''
    The FMA metadata is huge, for our experiment we only care about the "all_genres" tag.
    This method aims at creating a new - stripped down - CSV for FMA dataset
    Args:
        fma_meta_folder: the fma metadata folder path. Expected to find "genres.csv" and "tracks.csv"
        dst_dir: directory where the new dataset resides (where the new .csv will be dumped)
    '''
    from ast import literal_eval

    output = pd.DataFrame(columns = ['track_id','genres'])

    labels = pd.read_csv(os.path.join(fma_meta_folder,'tracks.csv'))
    labels = labels.loc[labels['set.1'].isin(['small', 'medium'])]
    label_genres = labels[['Unnamed: 0','track.9']]

    vocab = pd.read_csv(os.path.join(fma_meta_folder,'genres.csv'))

    for i, track in tqdm(label_genres.iterrows()):
        song = track['Unnamed: 0']
        try:
            song_genres = list(literal_eval(track['track.9']))
            literals = []
            if isinstance(song_genres, list):
                for genre_id in song_genres:
                    literals.append(vocab.loc[vocab['genre_id'] == genre_id]['title'].values[0])
                output.loc[len(output.index)] = [song,','.join(literals)]
        except Exception as e:
            print(e)
    
    output.to_csv(os.path.join(dst_dir,'fma_genres_ground_truth.csv'))

def build_librispeech_hq(librivox_dir, librispeech_dir, dst_dir, validate=True):
    '''
    Build librispeech segment files from a given LibriSpeech dataset directory.
    LibriSpeech is by default available in 16kHz. In order to use it in 44.1kHz, we need
    to build it from "source" (LibriVox). LibriSpeech is still needed as it contains the source
    metadata needed in order to reconstruct it from LibriVox
    Args:
        - librivox_dir: path of librivox to build from
        - librispeech_dir: path of librispeech to build from
        - dst_dir: dest. path of where librispeech-hq will be built
    '''
    def segment_to_dict(path):
        '''
        take a .seg text file and turn it into a dict of filename:timestamps pairs
        '''
        d = {}
        with open(path) as f:
            for line in f:
                key = line.split()[0]
                val = line.split()[1:]
                d[key] = val
        
        return d
    
    def utterance_to_dict(path):
        '''
        take a .map text file and turn it into a dict of filename1:filename2 pairs
        '''
        d = {}
        with open(path) as f:
            lines = f.readlines()
            # first two lines are special cases
            d['sentences only'] = False if 'False' in lines[0] else True
            for line in lines[2:]:
                key = line.split()[0]
                val = line.split()[1:]
                d[key] = val
            
        return d

    sr = 44100
    # Keep split names consistent across datasets
    split_format = {'train-clean-100':'train','dev-clean':'val','test-clean':'eval'}
    paths = glob(librivox_dir+'/**/*.mp3', recursive=True)
    speaker_per_set = {'test-clean':  [x.split('/')[-2] for x in glob(os.path.join(librispeech_dir,'test-clean','*/'))],
                       'dev-clean':   [x.split('/')[-2] for x in glob(os.path.join(librispeech_dir,'dev-clean','*/'))],
                       'train-clean-100': [x.split('/')[-2] for x in glob(os.path.join(librispeech_dir,'train-clean-100','*/'))]}
    for path in tqdm(paths):
        # First open text of all segments from orig path
        parent = path.split('/')[:-1]
        chapter_id = parent[-1]
        speaker_id = parent[-2]
        audio_filename = path.split('/')[-1].split('.')[0]
        current_split = [key for key, item in speaker_per_set.items() if speaker_id in item]

        if len(current_split) == 0:
            continue

        chapter_dir = Path(os.path.join(dst_dir,split_format[current_split[0]],speaker_id, chapter_id))
        chapter_dir.mkdir(parents=True, exist_ok= True)

        trans_filename = "{}-{}.trans.txt".format(speaker_id,chapter_id)
        # We need to copy the .trans file from src to dst
        shutil.copy(os.path.join(librispeech_dir, current_split[0], speaker_id, chapter_id, trans_filename),
                    chapter_dir)

        # Get the utterance dict
        utt_filename = "utterance_map.txt"
        utt_path = os.path.join('/'.join(parent[:-1]),utt_filename)
        utt_dict = utterance_to_dict(utt_path)

        # Get the segments dictionary corresponding to the "sentences only" value
        seg_ext = '.sents.seg.txt' if utt_dict["sentences only"] else '.seg.txt'
        utt_dict.pop('sentences only', None)
        seg_filename = "{}-{}{}".format(speaker_id, chapter_id, seg_ext)
        seg_path = os.path.join('/'.join(parent), seg_filename)
        seg_dict = segment_to_dict(seg_path)

        try:  
            # Load sound                                                   
            sound = AudioSegment.from_mp3(path)
            sound = sound.set_channels(1)
            data  = np.array(sound.get_array_of_samples())

            if sound.frame_rate != sr:
                data = resample(data, sound.frame_rate, sr)

            # 'key' is the new filename, 'item' is the segment to get the timestamp from (in 'seg_dict')
            for key, items in utt_dict.items():
                item = items[0]
                #check if we're still checking on the right audio file
                if '-{}-'.format(audio_filename) in key:
                    new_audio_path = os.path.join(chapter_dir,key+'.wav')
                    start_spl, end_spl = int(float(seg_dict[item][0])*sr), int(float(seg_dict[item][1])*sr)
                    new_data = data[start_spl:end_spl]
                    sf.write(new_audio_path, new_data, sr, 'PCM_16', format='WAV')
        except Exception as e: 
            print(e)
    
    if validate:
        validate_librispeech_hq(librispeech_dir, dst_dir)

def validate_librispeech_hq(librispeech_dir, librispeech_dir_new, t_thres = 0.05):

    all_files = ['/'.join(x.split('/')[3:]) for x in glob(librispeech_dir+'/**/*.wav', recursive=True)]

    for file in tqdm(all_files):
        
        low_a, l_sr = sf.read(os.path.join(librispeech_dir,file))
        high_a, h_sr = sf.read(os.path.join(librispeech_dir_new,file))

        if np.abs((len(low_a)/l_sr)-(len(high_a)/h_sr)) > t_thres:
            print('Warning, found file',file,'with time discrepancy larger than threshold: ',len(low_a)/l_sr,len(high_a)/h_sr)

def validate_audio(dirs):
    '''
    Given a directory of files, validate audio and check for errors.
    '''
    broken = []
    for path in dirs:
        paths   = glob(path+'/**/*.wav', recursive=True)
        for path in tqdm(paths):
            try:
                _, d = wavfile.read(path)
                if (np.isnan(d).any() or np.isposinf(d).any() or np.isneginf(d).any() or np.isinf(d).any()):
                    print(np.isnan(d).any(),
                        np.isposinf(d).any(),
                        np.isneginf(d).any(),
                        np.isinf(d).any())
                    broken.append(path)
            except Exception as e: 
                broken.append(path)
                print(e, path)

    np.savez('./dnr_broken_files.npz',broken)

def get_filelen_stats(path_dic, title=''):
    '''
    Given a directory of files, compute and plots length distribution in seconds
    Not used explicitly in the script, but good to have just in case.
    '''
    from collections import OrderedDict
    import matplotlib

    font = {'family' : 'normal','size': 22}
    matplotlib.rc('font', **font)
    matplotlib.rcParams['pdf.fonttype'] = 42
    matplotlib.rcParams['ps.fonttype'] = 42

    plt.cla()
    plt.clf()

    lengths = OrderedDict({'Librispeech':[],'FSD50k-Foreground':[],'FSD50k-Background':[],'FSD50k-Background':[],'FMA':[]})
    max_files = 1000
    for k, path in path_dic.items():
        print(k,path)
        paths = glob(path+'/**/*.wav', recursive=True)
        for path in tqdm(paths[:max_files]):
            try:
                data, sr = sf.read(path)
                lengths[k].append(data.shape[0] / sr)
            except Exception as e: 
                print(e)

    plt.figure(figsize=(12,6))
    plt.hist([x for x in lengths.values()], bins=50, label=list(lengths.keys()), alpha=0.5)
    plt.title('DNR per-class file length distribution')
    plt.legend(loc='best',bbox_to_anchor=(0.5, 0.5))
    plt.xlabel('File length (Sec)')
    plt.ylabel('Num. files')
    plt.xlim([0,32])
    plt.ylim([0,200])
    plt.tight_layout()
    plt.savefig('./img_'+title+'.pdf')

    return lengths

def main(args):

    # FSD50K Datasets: dev and eval set
    prepare_fsd50k(src_dir=args.fsd50k_path,
                   dst_dir=os.path.join(args.dest_dir,'FSD50K'))

    # Refactor the FSD50K dataset
    refactor_fsd50k(ground_truth_path=os.path.join(args.fsd50k_path,'FSD50K.ground_truth'),
                    root_path=os.path.join(args.dest_dir,'FSD50K'))
    
    # FMA Datasets: medium
    prepare_fma(src_dir=args.fma_path,
                dst_dir=os.path.join(args.dest_dir,'FMA'))

    # LibriSpeech: Build it off LibriVox
    build_librispeech_hq(librivox_dir=args.librivox_path, 
                         librispeech_dir=args.librispeech_path, 
                         dst_dir=os.path.join(args.dest_dir,'LibriSpeech'), 
                         validate=args.validate_audio)

    # Optionally, validate everything
    if args.validate_audio:
        validate_audio([os.path.join(args.dest_dir,'FMA'),
                        os.path.join(args.dest_dir,'FSD50K'),
                        os.path.join(args.dest_dir,'LibriSpeech')])

if __name__ == '__main__':
    '''
    Run this script in order to standardize all or parts of the datasets involved in the Divid and Remix (DNR) dataset.
    The script expects the three datasets to be found at a given directory (dest-dir). For more info on how to build and
    dataset-related specificities, please visit: <DNR website>

    Standardization involves:
        - Resampling based on the target sampling-rate
        - Summing all stereo files to mono (DNR is purely mono)
        - Rewriting files to a common format - WAV-PCM (some files are natively in .mp3 format (i.e. FMA))
        - Trimming the leading silences where necessary
        - Creating sets directory (val, eval, train)
        - FSD50K needs to be refactored in two classes: "background" and "foreground" prior to DNR building
    '''

    parser = argparse.ArgumentParser()
    parser.add_argument('--fsd50k-path', 
                        type=str, 
                        required=True, 
                        help='The path of the original FSD50k dataset')
    parser.add_argument('--fma-path', 
                        type=str, 
                        required=True, 
                        help='The path of the original FMA dataset')
    parser.add_argument('--librispeech-path', 
                        type=str, 
                        required=True, 
                        help='The path of the original LibriSpeech dataset')
    parser.add_argument('--librivox-path', 
                        type=str, 
                        required=True, 
                        help='The path of the original LibriVox dataset')
    parser.add_argument('--dest-dir', 
                        type=str, 
                        required=True, 
                        help='Destination directory where the datasets reside')
    parser.add_argument('--validate-audio', 
                        type=bool, 
                        required=False,
                        default=True,
                        help='Whether audio validation should be done at the end of the process or not')

    main(parser.parse_args())