p@24: """
p@24: This script creates lists of audio files contained in a folder.
p@24: 
p@24: Source code:
p@24: https://github.com/sidsig/ICASSP-MLP-Code/blob/master/make_lists.py
p@24: 
p@24: Modified by:
p@24: Paulo Chiliguano
p@24: MSc candidate Sound and Music Computing
p@24: Queen Mary University of London
p@24: 2015
p@24: 
p@24: References:
p@24:  - Sigtia, S., & Dixon, S. (2014, May). Improved music feature learning with deep neural 
p@24:    networks. In Acoustics, Speech and Signal Processing (ICASSP), 2014 IEEE International 
p@24:    Conference on (pp. 6959-6963). IEEE.
p@24: """
p@24: 
p@24: import numpy
p@24: import numpy.random as random
p@24: import os
p@24: import pickle
p@24: import sys
p@24: import utils as U
p@24: #import pdb
p@24: 
p@24: def read_file(filename):
p@24:     """
p@24:     Loads a file into a list
p@24:     """
p@24:     file_list=[l.strip() for l in open(filename,'r').readlines()]
p@24:     return file_list
p@24: 
p@24: def get_folds(filelist, n_folds):
p@24:     n_per_fold = len(filelist) / n_folds
p@24:     folds = []
p@24:     for i in range(n_folds-1):
p@24:         folds.append(filelist[i * n_per_fold: (i + 1) * n_per_fold])
p@24:     i = n_folds - 1
p@24:     folds.append(filelist[i * n_per_fold:])
p@24:     return folds
p@24: 
p@24: def generate_mirex_list(train_list, annotations):
p@24:     out_list = []
p@24:     for song in train_list:
p@24:         annot = annotations.get(song,None)
p@24:         if annot is None:
p@24:             print 'No annotations for song %s' % song
p@24:             continue
p@24:         assert(type('') == type(annot))
p@24:         out_list.append('%s\t%s\n' % (song,annot))
p@24: 
p@24:     return out_list
p@24:             
p@24: 
p@24: def make_file_list(gtzan_path, n_folds=5,):
p@24:     """
p@24:     Generates lists
p@24:     """
p@24:     audio_path = os.path.join(gtzan_path,'audio')
p@24:     out_path = os.path.join(gtzan_path,'lists')
p@24:     files_list = []
p@24:     for ext in ['.au', '.mp3', '.wav']:
p@24:         files = U.getFiles(audio_path, ext)
p@24:         files_list.extend(files)
p@24:     random.shuffle(files_list)
p@24:     
p@24:     if not os.path.exists(out_path):
p@24:         os.makedirs(out_path)
p@24:     
p@24:     audio_list_path = os.path.join(out_path, 'audio_files.txt')
p@24:     open(audio_list_path,'w').writelines(['%s\n' % f for f in files_list])
p@24:     
p@24:     annotations = get_annotations(files_list)
p@24: 
p@24:     ground_truth_path = os.path.join(out_path, 'ground_truth.txt')
p@24:     open(ground_truth_path,'w').writelines(generate_mirex_list(files_list, annotations))
p@24:     generate_ground_truth_pickle(ground_truth_path)
p@24: 
p@24:     folds = get_folds(files_list, n_folds=n_folds)
p@24:     
p@24:     ### Single fold for quick experiments
p@24:     create_fold(0, 1, folds, annotations, out_path)
p@24:     
p@24:     for n in range(n_folds):
p@24:         create_fold(n, n_folds, folds, annotations, out_path)
p@24: 
p@24: 
p@24: def create_fold(n, n_folds, folds, annotations, out_path):
p@24:     train_path = os.path.join(out_path, 'train_%i_of_%i.txt' % (n+1, n_folds))
p@24:     valid_path = os.path.join(out_path, 'valid_%i_of_%i.txt' % (n+1, n_folds))
p@24:     test_path = os.path.join(out_path, 'test_%i_of_%i.txt' % (n+1, n_folds))
p@24:     
p@24:     test_list = folds[n]
p@24:     train_list = []
p@24:     for m in range(len(folds)):
p@24:         if m != n:
p@24:             train_list.extend(folds[m])
p@24:     
p@24:     open(train_path,'w').writelines(generate_mirex_list(train_list, annotations))
p@24:     open(test_path,'w').writelines(generate_mirex_list(test_list, annotations))
p@24:     split_list_file(train_path, train_path, valid_path, ratio=0.8)
p@24:     
p@24: def split_list_file(input_file, out_file1, out_file2, ratio=0.8):
p@24:     input_list = open(input_file,'r').readlines()
p@24:     
p@24:     n = len(input_list)
p@24:     nsplit = int(n *ratio)
p@24:     
p@24:     list1 = input_list[:nsplit]
p@24:     list2 = input_list[nsplit:]
p@24:     
p@24:     open(out_file1, 'w').writelines(list1)
p@24:     open(out_file2, 'w').writelines(list2)
p@24: 
p@24: 
p@24: def get_annotation(filename):
p@24:     genre = os.path.split(U.parseFile(filename)[0])[-1]
p@24:     return genre
p@24: 
p@24: def get_annotations(files_list):
p@24:     annotations = {}
p@24:     for filename in files_list:
p@24:         annotations[filename] = get_annotation(filename)
p@24: 
p@24:     return annotations
p@24: 
p@24: def generate_ground_truth_pickle(gt_file):
p@24:     gt_path,_ = os.path.split(gt_file)
p@24:     tag_file = os.path.join(gt_path,'tags.txt')
p@24:     gt_pickle = os.path.join(gt_path,'ground_truth.pickle')
p@24:     
p@24:     lines = open(gt_file,'r').readlines()
p@24:     
p@24:     tag_set = set()
p@24:     for line in lines:
p@24:         filename,tag = line.strip().split('\t')
p@24:         tag_set.add(tag)
p@24:     tag_list = sorted(list(tag_set))
p@24:     open(tag_file,'w').writelines('\n'.join(tag_list + ['']))
p@24:     
p@24:     tag_dict = dict([(tag,i) for i,tag in enumerate(tag_list)])        
p@24:     n_tags = len(tag_dict)
p@24: 
p@24:     mp3_dict = {}
p@24:     for line in lines:
p@24:         filename,tag = line.strip().split('\t')
p@24:         tag_vector = mp3_dict.get(filename,numpy.zeros(n_tags))
p@24:         if tag != '':
p@24:             tag_vector[tag_dict[tag]] = 1.
p@24:         mp3_dict[filename] = tag_vector
p@24:     pickle.dump(mp3_dict,open(gt_pickle,'w'))
p@24: 
p@24: if __name__ == '__main__':
p@24:     if len(sys.argv) < 2:
p@24:         print 'Usage: python %s gtzan_path [n_folds=10]' % sys.argv[0]
p@24:         sys.exit()
p@24:     
p@24:     gtzan_path = os.path.abspath(sys.argv[1])
p@24:     if len(sys.argv) > 2:
p@24:         n_folds = int(sys.argv[2])
p@24:     else:
p@24:         n_folds = 10
p@24:         
p@24:     make_file_list(gtzan_path, n_folds)