audio.py
1 """ 2 Stream audio data over USB port (from bristlemouth!) 3 4 To run this script on MacOS, you must do the following: 5 * conda install portaudio 6 * pip3 install pyaudio 7 """ 8 import pyaudio 9 import serial 10 import wave 11 import sys 12 import argparse 13 from collections import namedtuple 14 import struct 15 import signal 16 17 sampling = True 18 def handler(signum, frame): 19 global wavfile 20 global stream 21 print("Exiting") 22 if wavfile: 23 wavfile.close() 24 25 stream.stop_stream() 26 stream.close() 27 28 sys.exit(0) 29 30 signal.signal(signal.SIGINT, handler) 31 32 33 AUDIO_HDR = namedtuple("AUDIO_HDR", "magic sample_rate num_samples sample_size flags") 34 AUDIO_STRUCT = "<LLHBB" 35 36 parser = argparse.ArgumentParser() 37 parser.add_argument("device", help="Serial port") 38 parser.add_argument("--filename", help="Save stream to wav file") 39 args, unknownargs = parser.parse_known_args() 40 41 ser = serial.Serial(args.device, 921600) 42 43 ser.flushInput() 44 45 print("Waiting for data") 46 47 # Wait for magic number 48 header_bytes = bytes([0x55, 0xB0, 0x10, 0xAD]) 49 50 ser.read_until(header_bytes) 51 52 # Get the first header 53 header_bytes += ser.read(struct.calcsize(AUDIO_STRUCT)-4) 54 header = AUDIO_HDR._make(struct.unpack(AUDIO_STRUCT, header_bytes)) 55 56 print("Starting stream!") 57 58 p = pyaudio.PyAudio() 59 stream = p.open(format=2, 60 channels=1, 61 rate=int(header.sample_rate/2), 62 output=True) 63 64 wavfile = None 65 if args.filename: 66 wavfile = wave.open(args.filename, "w") 67 wavfile.setnchannels(1) 68 wavfile.setsampwidth(header.sample_size) 69 wavfile.setframerate(header.sample_rate) 70 71 last_num_samples = header.num_samples 72 while True: 73 samples = ser.read(header.num_samples * header.sample_size) 74 stream.write(samples) 75 if wavfile: 76 wavfile.writeframes(samples) 77 78 if header.num_samples < last_num_samples: 79 print("Done") 80 break 81 82 # print(".") 83 # Read next header 84 header_bytes = ser.read(struct.calcsize(AUDIO_STRUCT)) 85 header = AUDIO_HDR._make(struct.unpack(AUDIO_STRUCT, header_bytes)) 86 87 if wavfile: 88 wavfile.close() 89 90 stream.stop_stream() 91 stream.close()