whack square wave midi player

midiplayer.py

import mido
import numpy as np
import scipy.signal
import scipy.io.wavfile
import simpleaudio as sa
import sys
import io
A = 440 # hertz
SAMPLE_RATE = 48000
wave_function = [np.sin, scipy.signal.square, scipy.signal.sawtooth][1]

round = int
is_stdin = False
def times_sr(n):
    return round(n * SAMPLE_RATE)

def parse_midi(fname):
    """converts the midi to a set of tuples representing notes"""
    if not is_stdin:
        midi = mido.MidiFile(fname)
    else:
        midi = mido.MidiFile(file=io.BytesIO(sys.stdin.buffer.read()))
    assert midi.type != 0, "oh god why not type 1 midis i quit"
    m = list(midi)
    total_num_notes = len([msg for msg in m if not msg.is_meta])
    elapsed = 0
    keyboard = [None] * 128
    out = []
    for idx, msg in enumerate(m):
        elapsed += msg.time
        if msg.is_meta:
            continue
        elif msg.type == "note_on":
            if keyboard[msg.note] is None:
                keyboard[msg.note] = (msg.velocity, elapsed)
            else:
                print("warning: two note_on events! ignoring")
                continue
        elif msg.type == "note_off":
            if keyboard[msg.note] is None:
                print("warning: note_off on an off note! ignoring")
                continue
            else:
                out.append((msg.note, keyboard[msg.note][0], keyboard[msg.note][1], elapsed - keyboard[msg.note][1]))
                note = out[-1]
                #print(note[0], note[1], note[2] * 48000, note[3] * 48000)
                #print(out[-1])
                keyboard[msg.note] = None
        print("Parsed", idx, "/", total_num_notes, "MIDI messages.", end="\r")
    print()
    return (out, midi)

def note_freq(note):
    # note[0] is the key/pitch. since a4 is mapped to key 69, we use that in our conversion.
    return A * 2 ** ((note[0] - 69) / 12)

def note_to_waveform(note):
    # note[3] is the note duration, note[1] is note velocity/amplitude
    t = np.linspace(0, note[3], round(note[3] * SAMPLE_RATE), False)
    return note[1] * wave_function(note_freq(note) * t * 2 * np.pi)
        
if __name__ == "__main__":
    if sys.argv[1] == "-":
        sys.argv[1] = "[stdin]"
        is_stdin = True
    print("Parsing MIDI", sys.argv[1])
    notes, midi = parse_midi(sys.argv[1])
    midi_length = midi.length
    #print("Parsed", len(notes), "notes.")
    total_num_notes = len(notes)
    notes.sort(key=lambda n: n[0], reverse=True)
    
    pitch = -1
    ts = 0
    audio = np.zeros(times_sr(midi_length))
    line = np.zeros(times_sr(midi_length))

    # iterate note by note, merging the sound line with the output audio as we go
    for idx, note in enumerate(notes, start=1):
        if note[0] != pitch:
            # add the note line to the audio
            audio += line 
            pitch = note[0]
            ts = 0
            line = np.zeros(times_sr(midi_length))
            #print(f"\n{pitch:>3}: ", end="")
        _, vel, note_ts, note_duration = note
        # insert the note waveform into the current line
        line[times_sr(note_ts):times_sr(note_ts) + times_sr(note_duration)] = note_to_waveform(note)
        #print(" " * (note_ts - ts) + "+" * note_duration, end="")
        ts = note_ts + note_duration
        print("Generated", idx, "/", total_num_notes, "notes.", end="\r")
    print()
    
    # normalize to 16 bit range
    audio *= 32767 / np.max(np.abs(audio))
    # convert to 16 bit data
    audio = audio.astype(np.int16)
    #print("Generated", len(notes), "waveforms")
    print("Writing audio to ", (sys.argv[1] if not is_stdin else "output") + ".wav")
    scipy.io.wavfile.write(sys.argv[1] + ".wav", SAMPLE_RATE, audio)
    print("Playing audio...")
    player = sa.play_buffer(audio, 1, 2, SAMPLE_RATE)
    player.wait_done()