gather_data.py

import os
import keyboard
import pyaudio, wave
import numpy.random as random


# Settings
chunk = 1024
channels = 2
sample_format = pyaudio.paInt16 
sample_rate = 44100
key_gap = 0.25
path = 'generating_dataset/data'


# Initializing audio interface
print('Initializing audio interface...')
p = pyaudio.PyAudio()
stream = p.open(
    format=sample_format,
    channels=channels,
    rate=sample_rate,
    frames_per_buffer=chunk,
    input=True
)


# Initializing keyboard handlers
print('Initializing keyboard handlers...')
def record_key(e):
    print('Pressed', e.name)
    keys.append((i, e.name))
keyboard.on_press(record_key)


# Start recording
print('Recording started')
print('Start pressing keys. Press Ctrl+q to stop recording.')

recording = True
def stop_recording():
    print('Stopping recording...')
    # Stop recording
    global recording
    recording = False
    # Remove ctrl+q keystrokes
    keys.pop()
    keys.pop()
# Watch for ctrl+q press
keyboard.add_hotkey('ctrl+q', stop_recording)


# Gather data
frames = []
keys = []

i = 0
while recording:
    data = stream.read(chunk)
    frames.append(data)
    i += 1


# Stop recording
stream.stop_stream()
stream.close()
p.terminate()
print('Finished recording')


# Save audio recording
print('Saving wav files...', end='')

def mkdir(key):
    try:
        os.mkdir(path + "/" + key)
    except Exception:
        pass

mkdir('')

start = int(0.33 * key_gap / chunk * sample_rate)
end = int(0.67 * key_gap / chunk * sample_rate)

def save_file(i, key):
    wf = wave.open(f'{path}/{key}/{i}{random.randint(99999)}.wav', 'wb')
    wf.setnchannels(channels)
    wf.setsampwidth(p.get_sample_size(sample_format))
    wf.setframerate(sample_rate)
    wf.writeframes(b''.join(frames[i - start:i + end]))
    wf.close()

for (i, key) in keys:
    mkdir(key)
    save_file(i, key)


print('Done!')
print('Exiting...')