drumExtractor/sampleExtraction.py at master · nathanshaw/drumExtractor · GitHub

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
"""
File containing functions for extracting samples, and naming them
"""
import wave
from scipy.io import wavfile
################

"""
----------------------------------
Proto-functions
----------------------------------
"""

def extractSamples(data, srate, onsets):
    """
    Function for extracting samples from data with a list of onsets
    """
    sampleList = []
    sampNum = 0
    for i in range(len(data)-1):
        if onsets[i] == 1:#if we found a peak
            start = i#set start to that sample
            end = i+(srate)
            found = 0
            #lets find the next onset
            for j in range(i+1, i + srate):#if there is another onset within a third of a second
                if j < len(data)-1:
                    if found == 0:
                        if onsets[j] == 1:
                            end = j + i
                            found = 9
            #load up the sample data into a NP array
            if(end - start > srate//3):
                sampleList.append(data[start:end])
                sampNum = sampNum + 1
    #print('----------------')
    #print(sampleList)
    print('----------------')
    print(sampNum, ' samples extracted')
    print('----------------')
    return sampleList

#def simpleSamples(samples, label=None):


def exportSamples(samples, label=None):
    """
    Function for Exporting Arrays of Audio data as labeled .wav files
    -----------------
    Variables
    -----------------
    samples = incomming array of audio data
    label   = incomming array with labels for incomming audio data
            The function assumes the following :
                0 = data is not percussive, function discards the sample
                1 = data is percussive but of unknown catagory
                2 = data is catagorized as a Kick Drum
                3 = data is catagorized as a Snare
                4 = data is catagorized as a Clap
    -----------------
    Returns :
    -----------------

    -----------------
    """
    tonalCount = 0
    unknownCount = 0
    kickCount = 0
    snareCount = 0
    clapCount = 0
    sampleName = 'unknown00'
    if (label == None):
        print("No label data given, will export all samples with unknown type")
        for i in range(len(samples)):
            sampleName = 'unknown' + str(i) + '.wav'
            print(sampleName)
            createWave(samples[i], sampleName)
        return 0
    else :
        print("Labeling Samples According to Label Data ... ")
        for i in range(len(samples)):
            currentLabel = label[i]
            sampleName, tonalCount, unknownCount, kickCount, snareCount, clapCount = nameMyFile(currentLabel, tonalCount, unknownCount, kickCount, snareCount, clapCount)
            createWave(samples[i], sampleName)
        return 0

def createWave(sample, fileName, srate = 44100):
    wavfile.write(fileName, srate, sample)
    return 0

def createWaveOld(sample, fileName, srate = 44100):
    wavFile = wave.open(fileName, 'wb')
    wavFile.setnchannels(1)
    wavFile.setframerate(srate)
    wavFile.setsampwidth(2*15-1)
    #wavFile.setparams(1, 2, srate, 0, 'NONE', 'not compressed')
    wavFileData = True
    for i in range(len(sample)):
        wavFileData = wavFileData + wave.struct.pack('h', int(sample[i])) #converts to binary
    wavFile.writeframes(wavFileData)
    wavFile.close()

def sampleLength(sample, srate):
    print("Sample length : ", len(sample)/srate, " Seconds")
    return 0

def nameMyFile(labelData, tonalCount, unknownCount, kickCount, snareCount, clapCount):
    if labelData.any() == 0:
        fileName = 'tonal' + str(tonalCount) + '.wav'
        tonalCount = tonalCount + 1
    if (labelData.any()== 1):
        fileName = 'unknown' + str(unknownCount) + '.wav'
        unknownCount = unknownCount + 1
    if (labelData.any() == 2):
        fileName = 'Kick' + str(kickCount) + '.wav'
        kickCount += 1
    if (labelData.any() == 3):
        fileName = 'Snare' + str(snareCount) + '.wav'
        snareCount += 1
    if (labelData.any() == 4):
        fileName = 'Clap' + str(clapCount) + '.wav'

    return fileName, tonalCount, unknownCount, kickCount, snareCount, clapCount