/ tools / scripts / misc / audio.py
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()