added more code

app.py

@@ -25,6 +25,41 @@ from torchvision import transforms
 import k_diffusion as K
 
 
+from .pianoroll import regroup_lines, img_file_2_midi_file, square_to_rect, rect_to_square
+from .square_to_rect import square_to_rect
+
+
+
+def infer_mask_from_init_img(img, mask_with='grey'):
+    "note, this works whether image is normalized on 0..1 or -1..1, but not 0..255"
+    assert mask_with in ['blue','white','grey']
+    "given an image with mask areas marked, extract the mask itself"
+    if not torch.is_tensor(img):
+        img = ToTensor()(img)
+    print("img.shape: ", img.shape)
+    # shape of mask should be img shape without the channel dimension
+    if len(img.shape) == 3:
+        mask = torch.zeros(img.shape[-2:])
+    elif len(img.shape) == 2:
+        mask = torch.zeros(img.shape)
+    print("mask.shape: ", mask.shape)
+    if mask_with == 'white':
+        mask[ (img[0,:,:]==1) & (img[1,:,:]==1) & (img[2,:,:]==1)] = 1
+    elif mask_with == 'blue':
+        mask[img[2,:,:]==1] = 1  # blue
+    if mask_with == 'grey':
+        mask[ (img[0,:,:] != 0) & (img[0,:,:]==img[1,:,:])  & (img[2,:,:]==img[1,:,:])] = 1
+    return mask*1.0
+
+
+def count_notes_in_mask(img, mask):
+    "counts the number of new notes in the mask"
+    img_t = ToTensor()(img)
+    new_notes = (mask * (img_t[1,:,:] > 0)).sum() # green channel
+    return new_notes.item()
+
+
+
 def greet(name):
     return "Hello " + name + "!!"
 

find_min_max_notes.py

@@ -0,0 +1,44 @@
+#!/usr/bin/env python3
+
+# this routine reads in many image files which are 128 pixels high
+# and finds the lowest and highest non-black pixels across all images
+import sys
+from PIL import Image
+from control_toys.data import fast_scandir
+from tqdm import tqdm
+
+if __name__ == "__main__":
+    if len(sys.argv) < 2:
+        print(f"Usage: {sys.argv[0]} <image-dirs>")
+        sys.exit(1)
+
+    img_dirs = sys.argv[1:]
+
+    img_files = []
+    for imgdir in img_dirs:
+        i_subdirs, imfile = fast_scandir(imgdir, ['png', 'jpg', 'jpeg'])
+        if imfile != []: img_files = img_files + imfile
+
+    # remove any files with the 'transpose' in the name 
+    img_files = [f for f in img_files if 'transpose' not in f]
+
+    if len(img_files) == 0:
+        print(f"No image files found in {img_dirs}")
+        sys.exit(1)
+
+    min_note = 128
+    max_note = 0
+    for img_file in tqdm(img_files):
+        img = Image.open(img_file)
+        img = img.convert('L') 
+        for y in range(img.size[-1]):
+            for x in range(img.size[0]):
+                val = img.getpixel((x,y))
+                if val > 0:
+                    if y < min_note: min_note = y
+                    if y > max_note: max_note = y
+                    break
+    print(f"min_note = {min_note}, max_note = {max_note}")
+    transpose_max = 12
+    print(f"So with transpostion of +/-{transpose_max}, the range is", min_note-transpose_max, "to", max_note+transpose_max)
+    sys.exit(0)

img2midi.py

@@ -0,0 +1,51 @@
+#! /usr/bin/env python3
+
+# This script takes piano roll image files and converts them to MIDI
+# It was intended for the P909 dataset
+# It will create a directory of images for each MIDI file, where each image is a frame of the MIDI file
+
+import os
+import sys
+import pretty_midi
+import argparse
+from multiprocessing import Pool, cpu_count
+from tqdm import tqdm
+from control_toys.data import fast_scandir
+from functools import partial
+from control_toys.pianoroll import img_file_2_midi_file
+
+if __name__ == '__main__':
+    p = argparse.ArgumentParser(description=__doc__,
+                            formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    # use --no-onsets to disable requirment of onsets
+    p.add_argument('--diff', default='', help='diff against this background image, new notes go in instrument 2')
+    p.add_argument('--onsets', action=argparse.BooleanOptionalAction, default=True)
+    p.add_argument('--separators', default=0, type=int, help='draw separators every this many pixels. (0=none)')
+    p.add_argument("img_dirs", nargs='+', help="directories containing image files")
+    p.add_argument("output_dir", help="output directory")
+    args = p.parse_args()
+    #print('args = ',args)
+
+    img_dirs = args.img_dirs
+    output_dir = args.output_dir
+
+
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
+    img_files = []
+    if os.path.isfile(img_dirs[0]):
+        img_files = img_dirs
+    else:
+        for imgdir in img_dirs:
+            i_subdirs, imfile = fast_scandir(imgdir, ['png', 'jpg', 'jpeg'])
+            if imfile != []: img_files = img_files + imfile
+    if len(img_files) == 0:
+        print(f"No image files found in {img_dirs}")
+        sys.exit(1)
+
+    process_one = partial(img_file_2_midi_file, output_dir=output_dir, require_onsets=args.onsets, 
+                          separators=args.separators, diff_img_file=args.diff)
+    cpus = cpu_count()
+    with Pool(cpus) as p:
+        list(tqdm(p.imap(process_one, img_files), total=len(img_files), desc='Processing image files'))

midi2img.py

@@ -0,0 +1,88 @@
+#! /usr/bin/env python3
+
+# This script takes a directory of MIDI files and converts them to images
+# It was intended for the P909 dataset
+# It will create a directory of images for each MIDI file, where each image is a frame of the MIDI file
+
+# if you want to use the chord extractor, you need to run chord extractor on all the files, and then output
+# a list of all the unique chord names to all_chords.txt e.g.
+# $ cat ~/datasets/jsb_chorales_midi/*/*_chords.txt | awk -F'\t' '{print $3}' | sort | uniq -c | awk '{print $2}' > all_chords.txt 
+#   cat midis/*_chords.txt | awk -F'\t' '{print $3}' | sort | uniq -c | awk '{print $2}' > all_chords.txt
+#   find midis -name '*_chords.txt' -exec cat {} + | awk -F'\t' '{print $3}' | sort | uniq -c | awk '{print $2}' > all_chords.txt
+
+import os
+import sys
+from multiprocessing import Pool, cpu_count, set_start_method
+from tqdm import tqdm
+from control_toys.data import fast_scandir
+from functools import partial
+import argparse
+from control_toys.pianoroll import midi_to_pr_img
+from control_toys.chords import simplify_chord, POSSIBLE_CHORDS
+
+def wrapper(args, midi_file, all_chords=None):
+    return midi_to_pr_img(midi_file, args.output_dir, show_chords=args.chords, all_chords=all_chords, 
+                          chord_names=args.chord_names, filter_mp=args.filter_mp, add_onsets=args.onsets,
+                          remove_leading_silence=(not args.silence))
+
+
+if __name__ == '__main__':
+    p = argparse.ArgumentParser(description=__doc__,
+                                formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    p.add_argument('-c','--chords', action='store_true', help='infer chords and add markers')
+    p.add_argument('--chord-names', action='store_true', help='Add text for chord names')
+    p.add_argument('--filter-mp', default=True, help='filter out non-piano, non-melody instruments')
+    #p.add_argument('--onsets', default=True, type=bool, help='add onset markers')
+    p.add_argument('--onsets', default=True, action=argparse.BooleanOptionalAction, help='Produce onset markers')   # either --onsets or --no-onsets, default is...? 
+    p.add_argument('--silence', default=True, action=argparse.BooleanOptionalAction, help='Leave silence at start of song (True) or remove it (False)')   
+    p.add_argument('--start-method', type=str, default='fork',
+                   choices=['fork', 'forkserver', 'spawn'],
+                   help='the multiprocessing start method')
+    p.add_argument('--simplify', action='store_true', help='Simplify chord types, e.g. remove 13s')
+    p.add_argument('--skip-versions', default=True, help='skip extra versions of the same song')
+    p.add_argument("midi_dirs", nargs='+', help="directories containing MIDI files")
+    p.add_argument("output_dir", help="output directory")
+    args = p.parse_args()
+    print("args = ",args)
+
+    set_start_method(args.start_method)
+    midi_dirs, output_dir = args.midi_dirs, args.output_dir
+
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
+    if os.path.isdir(midi_dirs[0]):
+        midi_files = []
+        for mdir in midi_dirs:
+            m_subdirs, mf = fast_scandir(mdir, ['mid', 'midi'])
+            if mf != []: midi_files = midi_files + mf
+    elif os.path.isfile(midi_dirs[0]):
+        midi_files = midi_dirs
+
+
+    if args.skip_versions: 
+        midi_files = [f for f in midi_files if '/versions/' not in f]
+    print("len(midi_files) = ",len(midi_files)) # just a check for debugging
+               
+    if args.chords: 
+        # TODO: this is janky af but for now...
+        # Get a list of all unique chords from a premade text file list of possible chords
+        # to make the file in bash, assuming you've already run extract_chords.py: (leave sort alphabetical, don't sort in order of freq)
+        # cat */*_chords.txt | awk -F'\t' '{print $3}' | sort | uniq -c | awk '{print $2}' > all_chords.txt
+        #with open('all_chords.txt') as f:
+        #    all_chords = f.read().splitlines()
+        # use possible chords as all chords
+        all_chords = POSSIBLE_CHORDS  # now we just generate these
+        if args.simplify:
+            all_chords = list(set([simplify_chord(c) for c in all_chords]))
+        print("len(all_chords) = ",len(all_chords))   
+        print("all_chords = ",all_chords) # just a check for debugging  
+    else:
+        all_chords = None
+
+    process_one = partial(wrapper, args, all_chords=all_chords)
+    num_cpus = cpu_count()
+    with Pool(num_cpus) as p:
+        list(tqdm(p.imap(process_one, midi_files), total=len(midi_files), desc='Processing MIDI files'))
+
+    print("Finished")

pianoroll.py

@@ -0,0 +1,634 @@
+#!/usr/bin/env python3
+
+"""
+various routines for converting midi files to piano roll images and back
+Unless otherwise noted: Author: Scott H. Hawley, Feb-March 2024
+"""
+
+import os
+import torch 
+import torchvision 
+import torchvision.transforms as transforms
+from PIL import Image, ImageOps, ImageDraw, ImageFont
+import numpy as np
+import pretty_midi
+import matplotlib.pyplot as plt
+from .chords import chord_num_to_color, simplify_chord, CHORD_BORDER
+from .utils import rect_to_square, square_to_rect
+
+ONSET_STYLE = 'new'  # 'old'=onset markers on pixels before notes, 'new'=onset markers are part of notes
+
+def plot_piano_roll(pr_array): 
+    plt.figure(figsize=(8, 8))
+    plt.imshow(np.flipud(pr_array), aspect='auto')
+    plt.show()
+
+
+
+def piano_roll_to_pretty_midi(piano_roll, fs=8, program=0):
+    # this routine copied from https://github.com/jsleep/pretty-midi/blob/ba7d01e5796fedf3ca0a3528e48b5242d9d2ccc3/examples/reverse_pianoroll.py
+    '''Convert a Piano Roll array into a PrettyMidi object
+     with a single instrument.
+
+    Parameters
+    ----------
+    piano_roll : np.ndarray, shape=(128,frames), dtype=int
+        Piano roll of one instrument
+    fs : int
+        Sampling frequency of the columns, i.e. each column is spaced apart
+        by ``1./fs`` seconds.
+    program : int
+        The program number of the instrument.
+
+    Returns
+    -------
+    midi_object : pretty_midi.PrettyMIDI
+        A pretty_midi.PrettyMIDI class instance describing
+        the piano roll.
+
+    '''
+    notes, frames = piano_roll.shape
+    #print("piano_roll.T[piano_roll.T != 0] = ",piano_roll.T[piano_roll.T != 0],flush=True)
+    pm = pretty_midi.PrettyMIDI()
+    instrument = pretty_midi.Instrument(program=program)
+
+    # pad 1 column of zeros so we can acknowledge inital and ending events
+    piano_roll = np.pad(piano_roll, [(0, 0), (1, 1)], 'constant')
+
+    # use changes in velocities to find note on / note off events
+    velocity_changes = np.nonzero(np.diff(piano_roll).T)
+
+    # keep track on velocities and note on times
+    prev_velocities = np.zeros(notes, dtype=int)
+    note_on_time = np.zeros(notes)
+
+    for time, note in zip(*velocity_changes):
+        # use time + 1 because of padding above
+        velocity = np.clip(piano_roll[note, time + 1], 0, 127)
+        #print("piano_roll[note, time + 1], velocity = ",piano_roll[note, time + 1], velocity,flush=True)
+        time = time / fs
+        if velocity > 0:
+            if prev_velocities[note] == 0:
+                note_on_time[note] = time
+                prev_velocities[note] = velocity
+        else:
+            pm_note = pretty_midi.Note(
+                velocity=prev_velocities[note],
+                pitch=note,
+                start=note_on_time[note],
+                end=time)
+            instrument.notes.append(pm_note)
+            prev_velocities[note] = 0
+    pm.instruments.append(instrument)
+    return pm
+
+
+
+#### beginning of code copied from midi2img.py
+
+def find_first_note_start(midi):
+    """find the start time of the first note in the midi file
+       used to help alignment to beats/bars
+    """
+    first_start = 10000.0
+    for instrument in midi.instruments:
+        for note in instrument.notes:
+            if note.start < first_start:
+                first_start = note.start
+    return first_start
+
+
+def get_piano_rolls(midi, fs, remove_leading_silence=True, add_onsets=True, debug=False):
+    """Converts a pretty_midi object to a piano roll for each instrument"""
+    duration = midi.get_end_time()   # find out duration of the midi file
+    n_frames = int(np.ceil(duration * fs))   # calculate the number of frames
+
+    # create a piano roll for each instrument
+    # TODO: currently this is only setup for POP909 dataset, need to generalize for other datasets
+    piano_rolls = {'PIANO':  np.zeros((128, n_frames)),
+                   'MELODY': np.zeros((128, n_frames)),
+                   'TOTAL':  np.zeros((128, n_frames))}
+    if remove_leading_silence:
+        first_start = find_first_note_start(midi)
+
+    for instrument in midi.instruments: 
+        name = instrument.name.upper()
+        if name in ['MELODY', 'PIANO']:
+            if debug: print(f"get_piano_rolls: instrument.name = {name}")
+            for note in instrument.notes:
+                if remove_leading_silence:
+                    note.start -= first_start
+                    note.end -= first_start
+                start = int(np.round(note.start * fs))  # quantize start time to nearest 16th note
+                dur = (note.end - note.start)*fs        # quantize duration  (Tip: don't separately quantize start & end; that can lead to "double-rounding" errors)
+                #end = int(np.round(note.end * fs)) 
+                end = start + int(np.round(dur))   # round means some notes will get held a bit too long, but "floor" would err on the side of extra staccatto notes which I don't want
+                if end==start: end = start+1  # make sure note is at least 1 pixel long
+                piano_rolls[name][note.pitch, start:end] = note.velocity ## value of piano roll array for these pixels will be the note velocity. end+1 so that "end" index gets covered
+                piano_rolls['TOTAL'][note.pitch, start:end] = note.velocity
+                #if note.velocity in [65,59,49,100]: print("note = ",note)
+
+                # extra fun: make sure all note onsets pop
+                piano_rolls[name][note.pitch, start-1] = 0
+                piano_rolls['TOTAL'][note.pitch, start-1] = 0
+
+    # if remove_leading_silence and add_onsets: # we need to add one pixel for the red onset dot at the start
+    #     for instrument in piano_rolls:
+    #         piano_rolls[instrument] = np.insert(piano_rolls[instrument], 0, 0, axis=1)
+
+    return piano_rolls
+
+
+def piano_roll_to_img(pr_frame, # this is an array of shape (128, n_frames)
+                      output_dir, midi_name, instrument, 
+                      start_col=None, add_onsets=True, chords=None, chord_names=False, debug=False,
+                      onset_style=ONSET_STYLE, # 'new' or 'old'
+                      ):
+    os.makedirs(f"{output_dir}/{midi_name}", exist_ok=True)
+    filename = f"{output_dir}/{midi_name}/{midi_name}_{instrument}.png"
+    if start_col is not None: 
+        filename = filename.replace(".png",f"_{str(start_col).zfill(5)}.png") 
+    #if debug: print("pr_frame.T[pr_frame.T != 0] = ",pr_frame.T[pr_frame.T != 0])
+    #scaling_factor = 65 / 18  # found empiracally by lots of checking
+    #pr_frame = np.round(pr_frame * scaling_factor).astype(np.uint8)
+
+    scale_factor = 2  # velocity only goes up to 127, but colors go up to 255
+    green_channel = np.clip(np.round(pr_frame*scale_factor), 0, 255).astype(np.uint8)
+    rgb_image = np.dstack((np.zeros_like(green_channel), green_channel, np.zeros_like(green_channel)))
+    img = Image.fromarray(rgb_image,'RGB')
+
+    if add_onsets:        # add little onset markers (red dots)
+        if onset_style=='old':
+            # any black pixel that has a green pixel to its right is an onset. color it red
+            # note that x any are flipped from what you'd think, e.g. "img.size =  (2352, 128)"
+            for y in range(img.size[-1]):
+                for x in range(img.size[0]-1):
+                    if img.getpixel((x,y)) == (0,0,0) and img.getpixel((x+1,y)) != (0,0,0):
+                        img.putpixel((x,y), (255,0,0)) 
+        elif onset_style=='new':
+            # New version: 
+            # any green pixel with a black pixel on its left becomes a red pixel. or if first pixel on row is green, make it red (matchinf velocity)
+            # Thus red pixel counts as both onset and first part of note, so shortest notes (16ths) will appear as only red with no green
+            #   btw this seems to agree w/ polyffusion's approach (??)
+            for y in range(img.size[-1]):
+                x = 0
+                pxl = img.getpixel((x,y))
+                r,g,b = pxl
+                if is_green(*pxl):
+                    img.putpixel((0,y), (g,0,0))  # make the first pixel of the note red, matching the green intensity
+                for x in range(1, img.size[0]):
+                    pxl = img.getpixel((x,y))
+                    r,g,b = pxl
+                    if is_green(*pxl) and is_black(*img.getpixel((x-1,y))):
+                        img.putpixel((x,y), (g,0,0))
+        else:
+            print(f"Error: Unrecognized onset_style = {onset_style}. Exiting.")
+            return
+
+    img = img.transpose(Image.FLIP_TOP_BOTTOM) # flip it vertically for display purposes
+
+
+    if chords is not None: # add the chord colors for each time as a rectangles at the top and bottom
+        if chord_names:
+            font_size = 7
+            try:
+                myFont = ImageFont.truetype("arial.ttf", 7) #mac
+            except:
+                myFont = ImageFont.load_default(size=font_size)
+
+
+        for c in chords:
+            color = chord_num_to_color(c['chord_num'])
+            img.paste(color, (int(c['start']), img.size[-1]-CHORD_BORDER, int(c['end']), img.size[-1]))
+            img.paste(color, (int(c['start']), 0, int(c['end']), CHORD_BORDER))
+            if chord_names:
+                chord_name = c['chord_name'].replace(':','')
+                if debug: print(f"chord_name = {chord_name}, chord_num = {c['chord_num']}")
+                xpos = int(c['start'])
+                I1 = ImageDraw.Draw(img)
+                I1.text((xpos, 0), chord_name, font=myFont, fill=(255, 255, 255))
+
+    if debug: print("img.size = ",img.size)
+    if 0 in img.size:
+        print(f"Error: img.size = {img.size}. Skipping this file.")
+        return
+    
+    # # just make sure all blue is gone: 
+    # img_array = np.array(img)
+    # img_array[:, :, 2] = 0
+    # img = Image.fromarray(img_array)
+    
+    img.save(filename)
+
+
+
+def check_for_melody_piano(midi: pretty_midi.PrettyMIDI, debug=False):
+    has_melody, has_piano = False, False
+    if debug:
+        print("check_for_melody_piano: midi.instruments = ",midi.instruments)
+    for i, instrument in enumerate(midi.instruments):
+        if debug: print(f"check_for_melody_piano: instrument = [{instrument.name.upper()}]")
+        if instrument.name.upper() == 'MELODY': has_melody = True
+        if instrument.name.upper() == 'PIANO':  has_piano = True
+    # if theres only one instrument with no name, name it PIANO
+    if len(midi.instruments) == 1 and midi.instruments[0].name == '':
+        has_piano = True
+        midi.instruments[0].name = 'PIANO'
+    return has_melody, has_piano
+
+
+
+def midi_to_pr_img(midi_file, output_dir, 
+                   show_chords=None, # to show chords or not
+                   all_chords=None,  # list of all possible chords
+                   add_onsets=True,  # add red dots for note onsets
+                   chord_names=False, # to show chord names or not
+                   filter_mp=True,   # filter midi & piano
+                   remove_leading_silence=True, # remove silence at start of song
+                   simplify_chords=False, # simplify chord names
+                   debug=False,):    # show debugging info 
+    """Converts a MIDI file to a piano roll image"""
+    if debug: print(f"midi_to_pr_img: Processing {midi_file}")
+    if '/versions/' in midi_file and args.skip_versions: return
+    midi = pretty_midi.PrettyMIDI(midi_file)
+
+    if not check_for_melody_piano(midi):
+        print(f"Not ok: File {midi_file} does not have melody and piano. Skipping")
+        return
+    else:
+        if debug: print(f"Ok: File {midi_file} has melody and piano. Processing") 
+
+    ### Normalize tempo to 120bpm
+    tempo_changes = midi.get_tempo_changes()
+    start_tempo = tempo_changes[1][0]
+    bps = start_tempo / 60.0
+    fs = bps * 4.0 * 2
+    if debug: print("start_tempo, fs = ", start_tempo, fs)
+
+    chords=None
+    if show_chords and all_chords is not None:
+        # read the chord timing file, but note that those times have not yet been normalized to 120bpm 
+        # this file has column-separated format "start_time end_time chord"
+        chords_path = midi_file.replace('.mid', '_chords.txt')
+        with open(chords_path) as f:
+            chords = f.read().splitlines()
+        # split each line of text into a dict 3 values {'start':, 'end':, 'chord':}:
+        chords = [dict(zip(['start', 'end', 'chord'], c.split('\t'))) for c in chords]
+
+        for c in chords:
+            c['start'] = int(np.floor(float(c['start']) * fs))
+            c['end'] = int(np.ceil(float(c['end']) * fs))
+            c['chord_name'] = simplify_chord(c['chord']) if simplify_chords else c['chord']     
+            c['chord_num'] = all_chords.index(c['chord_name'])              
+
+    if filter_mp:  # remove non-piano, non-melody instruments
+        midi.instruments = [instrument for instrument in midi.instruments if instrument.name.upper() in ['MELODY', 'PIANO']]
+  
+    piano_rolls = get_piano_rolls(midi, fs, remove_leading_silence=remove_leading_silence, add_onsets=add_onsets)
+    if debug:
+        for p in piano_rolls.keys():
+            print(f"p  {p}.shape =",piano_rolls[p].shape)
+            #print(f"piano_rolls[{p}][piano_rolls[p] != 0] = ",piano_rolls[p][piano_rolls[p] != 0])
+ 
+    midi_name = os.path.basename(midi_file).split('.')[0] # get the midi filename w/o parent dirs or file extension
+
+    for instrument in piano_rolls: # save each instrument's piano roll as a single image
+        if debug: print("saving piano roll for ",instrument)
+        piano_roll_to_img(piano_rolls[instrument], output_dir, midi_name, instrument, chords=chords, chord_names=chord_names, 
+                          add_onsets=add_onsets, debug=debug)
+
+    return
+
+#### end of code copied from midi2img.py
+
+
+#### below code originally in img2midi.py
+
+def blockout_topbottom_arr(arr, border=CHORD_BORDER):
+    "set the top and bottom border pixels to black"
+    arr2 = arr.copy()
+    arr2[:border, :] = 0
+    arr2[-border:, :] = 0
+    return arr2
+
+def filter_by_velocity(midi, thresh=20):
+    "filter out notes with velocities below a certain threshold"
+    for instrument in midi.instruments:
+        notes = [note for note in instrument.notes if note.velocity > thresh]
+        instrument.notes = notes
+    return midi
+
+def img2midi(img, draw_sep=512, debug=False):
+    # operates on a single image
+    # flip the image vertically because numpy and PIL have different ideas of what the first row is
+    # if image vertical dimension is more than 128, then cut it into strips of 128 and concatenate them horizontally
+    if debug: print(f"img2midi: img.size = {img.size}")
+    if img.size[1] > 128:
+        arr = np.concatenate([np.array(img.crop((0, i, img.size[0], i+128))) for i in range(0, img.size[1], 128)], axis=1)
+    else:
+        arr = np.array(img)
+    if debug: print("0: arr.T[arr.T != 0] = ",arr.T[arr.T != 0])
+    arr = blockout_topbottom_arr(arr)
+
+    scale_factor = 0.5 # rgb down to velocity values
+    piano_roll_array = np.array(arr*scale_factor, dtype=np.int32)
+    piano_roll_array = np.flip(piano_roll_array, axis=0) # numpy as PIL are upside down relative to each other
+    if debug: 
+        print(f"piano_roll_array.shape = {piano_roll_array.shape}, piano_roll_array.dtype = {piano_roll_array.dtype}")
+        print("1: piano_roll_array[piano_roll_array != 0] = ",piano_roll_array[piano_roll_array != 0])
+    
+    # draw a vertical line every 128/256/512 pixels
+    if draw_sep > 0:
+        line_every = draw_sep
+        for i in range(0, piano_roll_array.shape[-1], line_every):
+            if i>0: piano_roll_array[35:-35,i] = 30
+    piano_roll_array = np.clip(piano_roll_array, 0, 127) # make sure velocities aren't out of bounds
+    midi = piano_roll_to_pretty_midi(piano_roll_array)
+    midi = filter_by_velocity(midi)
+    return midi
+
+
+def flip_bottom_half_and_attach(sub_img):
+    "takes one 256x256 and returns on 512x128 image with the bottom half reversed and attached on the right"
+    h, w = sub_img.size
+    new_img = Image.new(sub_img.mode, (w*2, h//2))
+    new_img.paste(sub_img.crop((0, 0, w, h//2)), (0, 0))
+    new_img.paste(sub_img.crop((0, h//2, w, h)).transpose(Image.FLIP_LEFT_RIGHT), (w, 0))
+    return new_img 
+
+
+def square_to_rect(img):
+    #"""just an alias for flip_bottom_half_and_attach"""
+    return flip_bottom_half_and_attach(img)
+
+def rect_to_square(img):
+    "takes a 512x128 image and returns a 256x256 image with the bottom half reversed"
+    w, h = img.size
+    new_img = Image.new(img.mode, (w//2, h*2))
+    new_img.paste(img.crop((0, 0, w//2, h)), (0, 0))
+    new_img.paste(img.crop((w//2, 0, w, h)).transpose(Image.FLIP_LEFT_RIGHT), (0, h))
+    return new_img
+
+def regroup_lines(img, debug=False): 
+    """
+    large images come in as an 8x8 grid of 256x256 images, in which the bottom half of each 256x256 is horizontally backwards
+    we will rebuild this grid by first flipping the bottom half of each 256x256 image
+    """
+    img2 = Image.new('RGB', img.size)
+    if debug: print(f"regroup_lines: img.size = {img.size}")
+
+    if img.size[0] == 256:
+        img2 = Image.new('RGB', (512,128))
+    elif img.size[0] != 2048:
+        if debug: print("regroup_lines: unexpected image size, returning image unchanged")
+        return  img # no op, hope all's well
+    imnum = 0
+    for row in range(0, img.size[0], 256):
+        for col in range(0, img.size[1], 256):
+            imnum += 1
+            sub_img = img.crop((col, row, col+256, row+256))
+            sub_img = square_to_rect(sub_img)
+            paste_x, paste_y = (imnum-1) % 4 * 512, (imnum-1) // 4 * 128
+            if debug: print(f"imnum = {imnum}, paste_x = {paste_x}, paste_y = {paste_y}")
+            img2.paste(sub_img, (paste_x, paste_y))
+    if debug: img2.show()
+    return img2
+
+
+def is_red(r,g,b, thresh=20, debug=False):
+    result = r > thresh and g < thresh and b < thresh
+    if debug: print("is_red: r,g,b = ",r,g,b,", result = ",result)
+    return result
+
+def is_green(r,g,b, thresh=20):
+    return r < thresh and g > thresh and b < thresh
+
+def is_black(r,g,b, thresh=20):
+    return r < thresh and g < thresh and b < thresh
+
+def filter_redgreen(img:Image, 
+                    require_onsets=True, # only keep green lines that start with a red pixel on the left
+                    thresh=20, # minimum amount of red or green to count
+                    onset_style=ONSET_STYLE, # 'new' or 'old
+                    debug=False):
+    # filter: only keep black points, and green lines that start with a red pixel on the left.
+    # i.e. only green points that have red or green to their left are valid notes
+    # intended for img2midi 
+    img.save('rgfilter_in.png')
+    img2 = img.copy()
+    if debug: print("img.size = ",img.size,", require_onsets = ",require_onsets," (not require_onsets) =",(not require_onsets)," thresh = ",thresh)
+    w, h = img.size
+    for y in range(CHORD_BORDER,h-CHORD_BORDER):
+        note_on = False
+        for x in range(w):  # scan from right to left
+            r,g,b = img2.getpixel((x,y))               # pixel under consideration
+            if debug and (r,g,b)!=(0,0,0): print(f"x, y: {x}, {y}:  r, g, b = {r},{g},{b}, note_on = {note_on}, is_red = {is_red(r,g,b, thresh)}, is_green = {is_green(r,g,b, thresh)}")
+            if is_red(r,g,b, thresh): 
+                note_on = True   
+                if onset_style == 'new': # keep the note but change the red to green
+                    img2.putpixel((x,y), (0,r,0))
+            elif is_green(r,g,b, thresh) and require_onsets and note_on:
+                img2.putpixel((x,y), (r,g,b)) # keep the note
+            elif is_green(r,g,b, thresh) and (not require_onsets):
+                img2.putpixel((x,y), (r,g,b)) # keep the note
+                note_on = True 
+            else:
+                note_on = False
+                img2.putpixel((x,y), (0,0,0))  # zero it out
+    
+    img2.save('rgfilter_out.png') # debugging always on here
+    return img2
+
+def arr_check(img, tag=''):
+        img = img.convert("RGB")
+        arr = np.array(img)[:,:,1]
+        print(tag,": arr.shape = ",arr.shape, flush=True)
+        print(tag,": arr.T[arr.T != 0] = ",arr.T[arr.T != 0], flush=True)
+
+def img2midi_multi(img, require_onsets=True, separators=512, debug=False):
+    "can operate on a grid of images"
+    img = img.convert('RGB') 
+    img = regroup_lines(img)
+    img = filter_redgreen(img, require_onsets=require_onsets)
+    #img = img.convert('L') # convert to grayscale
+    red_arr = np.array(img.split()[0])
+    green_arr = np.array(img.split()[1])
+    combined_arr = red_arr + green_arr
+    if debug: arr_check(img, '1')
+    img = Image.fromarray(combined_arr, mode="L")
+    if debug: arr_check(img, '2')
+    return img2midi(img, draw_sep=separators)
+
+def infer_mask_from_init_img(img, mask_with='grey', debug=True):
+    "note, this works whether image is normalized on 0..1 or -1..1, but not 0..255"
+    assert mask_with in ['blue','white','grey']
+    "given an image with mask areas marked, extract the mask itself"
+    img = np.array(img)
+    mask = np.zeros(img.shape[:2])
+    if debug: print("infer: img.shape, mask.shape = ",img.shape, mask.shape)
+    if mask_with == 'white':
+        mask[ (img[0,:,:]==1) & (img[1,:,:]==1) & (img[2,:,:]==1)] = 1
+    elif mask_with == 'blue':
+        mask[img[2,:,:]==1] = 1  # blue
+    if mask_with == 'grey':
+        mask[ (img[:,:,0] != 0) & (img[:,:,0] > -1) & (img[:,:,0]==img[:,:,1])  & (img[:,:,2]==img[:,:,1])] = 1
+    return mask*1.0
+
+def img_file_2_midi_file(img_file, output_dir='', require_onsets=True, separators=512, 
+                         diff_img_file='', debug=True):
+    "Converts an image file to a midi file"
+    if debug: print(f"Processing {img_file}", flush=True)
+    img = Image.open(img_file)
+    if debug: arr_check(img, '0')
+    midi = img2midi_multi(img, require_onsets=require_onsets, separators=separators)
+
+    if diff_img_file != '': # put new notes on new instrument. 
+        bg_img = rect_to_square(Image.open(diff_img_file))
+        mask = infer_mask_from_init_img(bg_img, mask_with='grey')
+        # tile mask to 3 color channels
+        mask = np.stack([mask]*3, axis=-1)
+        if debug: 
+            print("mask.shape = ",mask.shape, flush=True)
+            print("mask.min(), mask.max(), mask.sum() = ",mask.min(), mask.max(), mask.sum(), flush=True)
+        if bg_img.size[0] > img.size[0]:
+            bg_img = rect_to_square(bg_img)
+        bg_midi = img2midi_multi(bg_img, require_onsets=require_onsets, separators=separators)
+        # grab just the pixels in the mask of img
+        arr = np.array(img) 
+        if debug: 
+            print("arr.shape, mask.shape = ",arr.shape, mask.shape, flush=True)
+            print("arr.min(), arr.max(), arr.sum() = ",arr.min(), arr.max(), arr.sum(), flush=True)
+        new_arr = np.zeros_like(arr)
+        new_arr[mask > 0] = 1*255# arr[mask > 0]
+        new_arr = np.where(mask>0, arr, 0)
+        # new_arr[:,:,0] = arr[:,:,0] * mask    
+        # new_arr[:,:,1] = arr[:,:,1] * mask    
+        # new_arr[:,:,2] = arr[:,:,2] * mask    
+        new_img = Image.fromarray(new_arr, 'RGB')
+        square_to_rect(new_img).save('new_img.png')
+        new_midi = img2midi_multi(new_img, require_onsets=require_onsets, separators=separators)
+        bg_midi.instruments.append(new_midi.instruments[0])
+        midi = bg_midi 
+
+    midi_file = os.path.basename(img_file).replace('.png', '.mid')
+    if output_dir is not None and output_dir != '':
+        midi_file = os.path.join(output_dir, midi_file)
+    midi.write(midi_file)
+    return midi_file
+
+
+
+#### end of code copied from img2midi.py
+
+
+
+### dataset routines, called from train.h
+
+class RandomVerticalShift(torch.nn.Module):
+    """
+    Update: UNUSED.  Instead we do all transposing as pre-processing to facilitate chord detection.
+    Randomly shift the image vertically by up to max_shift pixels, which correspond to semitones.
+    """
+    def __init__(self, max_shift=12):
+        super().__init__()
+        self.max_shift = max_shift
+
+    def __call__(self, img):
+        shift = torch.randint(-self.max_shift, self.max_shift, (1,))
+        return self.vertical_shift(img, shift.item())
+
+    def vertical_shift(self, img, shift):
+        img = ImageOps.exif_transpose(img)  # Handle EXIF Orientation
+        img = img.transform(img.size, Image.AFFINE, (1, 0, 0, 0, 1, shift), fillcolor=0)
+        return img
+
+
+
+class RandomBarCrop(torch.nn.Module):
+    """
+    Given a PIL image of a piano roll (non-square!), do random cropping the level of measures, i.e. bars, i.e. 16 16th-note pixels
+    NOTE: might be nice if piano roll images have initial silence pre-removed -- assuming the first note is supposed to start on the first beat
+    """
+    def __init__(self, bar_length=16, window_length=512):
+        super().__init__()
+        self.bl = bar_length  # in 16th notes (16 pixels)
+        self.wl = window_length  # in pixels
+        self.bic = self.wl // self.bl  # bars in crop
+
+    def __call__(self, img: Image, debug=False):
+        bars_in_image = img.size[0] // self.bl  # number of bars in full image
+        if self.bic >= bars_in_image: # pad horizontal end of image with zeros if needed
+            pad = self.wl - img.size[0] + 1
+            img = ImageOps.expand(img, (0, 0, pad, 0), fill=0)
+            bars_in_image = img.size[0] // self.bl
+        try:
+            start_ind = torch.randint(0, bars_in_image - self.bic+1, (1,)).item() # start index of crop
+        except Exception as e:
+            print(f"***MY ERROR: {e}.  bars_in_image = {bars_in_image}, self.bic = {self.bic}")
+            assert False
+        start_pixel = start_ind * self.bl  # start pixel of crop
+        new_img = img.crop((start_pixel, 0, start_pixel + self.wl, img.size[1]))
+        assert new_img.size[0] == self.wl and new_img.size[1]==128, f"ERROR: new_img.size = {new_img.size}, self.wl = {self.wl}"
+        return new_img
+
+
+
+class StackPianoRollsImage(torch.nn.Module):
+    """
+    Given a PIL image of a piano roll, cut in in half horizontally, 
+    stack the two halves, with the lower half mirrored horzontally.
+    """
+    def __init__(self, final_size=(256, 256), max_shift=13):
+        super().__init__()
+        self.final_size = final_size
+
+    def __call__(self, img: Image, debug=False):
+        if img.size[0] <= 128 and img.size[1] <= 128:
+            return img # don't stack small images
+        # image dimensions are likely 512x128. I want 256x256 output
+        half_width = img.size[0] // 2
+        #make a new image with dimensions 256x256, with the same color mode as img
+        new_img = Image.new(img.mode, self.final_size)
+        # paste the first half of the image into the top half of the new image
+        first_half = img.crop((0, 0, half_width, img.size[1]))
+        new_img.paste(first_half, (0, 0))
+        # paste the second half of the image into the bottom half of the new image, but flipped horizontally
+        next_half = img.crop((half_width, 0, 2*half_width, img.size[1])) 
+        next_half = ImageOps.mirror(next_half)
+        new_img.paste(next_half, (0, img.size[1]))
+        return new_img
+
+
+
+class StackPianoRollsTensor(torch.nn.Module):   
+    """
+    Tensor version of StackPianoRollsImage.  Unused, i think. 
+    Given a torch tensor of a piano roll, cut in in half horizontally, stack the two halves
+    but have the bottom half mirrored horzontally.
+    """
+    def __init__(self):
+        super().__init__()
+
+    def __call__(self, img: torch.Tensor):
+        if img.shape[0] <= 128 and img.shape[1] <= 128:
+            return img # don't stack small images
+        img = img.permute(1, 2, 0)
+        half_width = img.shape[0] // 2
+        img = torch.cat([img[:half_width], img[half_width:][::-1]], dim=0)
+        img = img.permute(2, 0, 1)
+        return img
+    
+
+
+
+if __name__ == '__main__':
+    import sys 
+
+    # testing for  the StackPianoRollsImage class
+    filename = sys.argv[-1]
+    print("filename = ", filename)
+    img = Image.open(filename)
+    img = transforms.RandomCrop((128, 512))(img) # randomly crop it to 128x512
+    img = StackPianoRollsImage()(img, debug=True)
+    img.show()
+

rect_to_square.py

@@ -0,0 +1,26 @@
+#!/usr/bin/env python3
+
+# reads in a rectangular image, cuts it in half length-wise, and attaches the right half to the bottom, reversed
+
+import sys 
+from PIL import Image
+
+
+def rect_to_square(img):
+    # get width and height of img
+    w, h = img.size
+    new_img = Image.new(img.mode, (w//2, h*2))
+    new_img.paste(img.crop((0, 0, w, h)), (0, 0))
+    new_img.paste(img.crop((0, 0, w, h)).transpose(Image.FLIP_LEFT_RIGHT), (0, h))
+    return new_img
+
+
+if __name__ == '__main__':
+    img_filename = sys.argv[1]
+    in_img = Image.open(img_filename)
+    print("Input image dimenions: ", in_img.size)
+    out_img = rect_to_square(in_img)
+    print("Output image dimenions: ", out_img.size)
+    out_filename = img_filename.replace('.png', '_square.png')
+    print("Saving to ", out_filename)
+    out_img.save(out_filename)

square_to_rect.py

@@ -0,0 +1,31 @@
+#!/usr/bin/env python3
+
+# reads in a square image and outputs a rectangular image
+
+import sys 
+from PIL import Image
+
+def flip_bottom_half_and_attach(img):
+    "takes one 256x256 and returns on 512x128 image with the bottom half reversed and attached on the right"
+    #w, h = 256, 256
+    h, w = img.size
+    #assert sub_img.size[0] == h and sub_img.size[1] == w
+    new_img = Image.new(img.mode, (w*2, h//2))
+    new_img.paste(img.crop((0, 0, w, h//2)), (0, 0))
+    new_img.paste(img.crop((0, h//2, w, h)).transpose(Image.FLIP_LEFT_RIGHT), (w, 0))
+    return new_img 
+
+def square_to_rect(img):
+    # just an alias for flip_bottom_half_and_attach
+    return flip_bottom_half_and_attach(img)
+
+
+if __name__ == '__main__':
+    img_filename = sys.argv[1]
+    square_img = Image.open(img_filename)
+    print("Input image dimenions: ", square_img.size)
+    rect_img = square_to_rect(square_img)
+    print("Output image dimenions: ", rect_img.size)
+    out_filename = img_filename.replace('.png', '_rect.png')
+    print("Saving to ", out_filename)
+    rect_img.save(out_filename)