update MIDITokenizer

midi_tokenizer.py

@@ -200,7 +200,7 @@ class MIDITokenizerV1:
                     note_tracks = channel_note_tracks[c]
                     if len(note_tracks) != 0 and track_idx not in note_tracks:
                         track_idx = channel_note_tracks[c][0]
-                    new_track_idx = tr_map.setdefault(track_idx, next(iter(tr_map.values())))
+                    new_track_idx = tr_map[track_idx]
                     event[3] = new_track_idx
                     if name == "patch_change" and event[4] not in patch_channels:
                         patch_channels.append(event[4])
@@ -235,7 +235,9 @@ class MIDITokenizerV1:
             else:
                 if event[0] == "note":
                     notes_in_setup = True
-                key = tuple(event[3:-1])
+                    key = tuple([event[0]] + event[3:-2])
+                else:
+                    key = tuple([event[0]] + event[3:-1])
             setup_events[key] = new_event
 
         last_t1 = 0
@@ -551,6 +553,8 @@ class MIDITokenizerV2:
     def detect_key_signature(key_hist, threshold=0.7):
         if len(key_hist) != 12:
             return None
+        if sum(key_hist) == 0:
+            return None
         p = sum(sorted(key_hist, reverse=True)[:7]) / sum(key_hist)
         if p < threshold:
             return None
@@ -590,7 +594,7 @@ class MIDITokenizerV2:
         empty_channels = [True] * 16
         channel_note_tracks = {i: list() for i in range(16)}
         note_key_hist = [0]*12
-        key_sig_num = 0
+        key_sigs = []
         track_to_channels = {}
         for track_idx, track in enumerate(midi_score[1:129]):
             last_notes = {}
@@ -661,11 +665,11 @@ class MIDITokenizerV2:
                     sf, mi = event[2:]
                     if not (-7 <= sf <= 7 and 0 <= mi <= 1):  # invalid
                         continue
-                    key_sig_num += 1
                     sf += 7
                     new_event += [sf, mi]
+                    key_sigs.append(new_event)
 
-                if name == "note":
+                if name in ["note", "time_signature", "key_signature"]:
                     key = tuple(new_event[:-2])
                 else:
                     key = tuple(new_event[:-1])
@@ -731,7 +735,9 @@ class MIDITokenizerV2:
 
             empty_channels = [channels_map[c] for c in empty_channels]
             track_idx_dict = {}
+            key_sigs = []
             key_signature_to_add = []
+            key_signature_to_remove = []
             for event in event_list:
                 name = event[0]
                 track_idx = event[3]
@@ -748,14 +754,23 @@ class MIDITokenizerV2:
                     for c, tr_map in track_idx_map.items():
                         if track_idx in tr_map:
                             new_track_idx = tr_map[track_idx]
+                            c = channels_map[c]
                             new_channel_track_idx = (c, new_track_idx)
+                            if new_track_idx == 0:
+                                continue
                             if new_channel_track_idx not in new_channel_track_idxs:
                                 new_channel_track_idxs.append(new_channel_track_idx)
+
                     if len(new_channel_track_idxs) == 0:
-                        event[3] = 0
+                        if event[3] == 0: # keep key_signature on track 0 (meta)
+                            key_sigs.append(event)
+                            continue
+                        event[3] = -1 # avoid remove same event
+                        key_signature_to_remove.append(event) # empty track
                         continue
                     c, nt = new_channel_track_idxs[0]
                     event[3] = nt
+                    key_sigs.append(event)
                     if c == 9:
                         event[4] = 7 # sf=0
                     for c, nt in new_channel_track_idxs[1:]:
@@ -763,6 +778,7 @@ class MIDITokenizerV2:
                         new_event[3] = nt
                         if c == 9:
                             new_event[4] = 7  # sf=0
+                        key_sigs.append(new_event)
                         key_signature_to_add.append(new_event)
                 elif name == "control_change" or name == "patch_change":
                     c = event[4]
@@ -772,10 +788,12 @@ class MIDITokenizerV2:
                     note_tracks = channel_note_tracks[c]
                     if len(note_tracks) != 0 and track_idx not in note_tracks:
                         track_idx = channel_note_tracks[c][0]
-                    new_track_idx = tr_map.setdefault(track_idx, next(iter(tr_map.values())))
+                    new_track_idx = tr_map[track_idx]
                     event[3] = new_track_idx
                     if name == "patch_change" and event[4] not in patch_channels:
                         patch_channels.append(event[4])
+            for key_sig in key_signature_to_remove:
+                event_list.remove(key_sig)
             event_list += key_signature_to_add
             track_to_channels ={}
             for c, tr_map in track_idx_map.items():
@@ -793,16 +811,31 @@ class MIDITokenizerV2:
                 if c not in patch_channels and c in track_idx_dict:
                     event_list.append(["patch_change", 0, 0, track_idx_dict[c], c, 0])
 
-        if key_sig_num == 0:
-            # detect key signature.
+        if len(key_sigs) == 0 or all([key_sig[4]==7 for key_sig in key_sigs]):
+            # detect key signature or fix the default key signature
             root_key = self.detect_key_signature(note_key_hist)
             if root_key is not None:
                 sf = self.key2sf(root_key, 0)
                 # print("detect_key_signature",sf)
-                for tr, cs in track_to_channels.items():
-                    if remap_track_channel and tr == 0:
-                        continue
-                    event_list.append(["key_signature", 0, 0, tr, (0 if (len(cs) == 1 and cs[0] == 9) else sf) + 7, 0])
+                if len(key_sigs) == 0:
+                    for tr, cs in track_to_channels.items():
+                        if remap_track_channel and tr == 0:
+                            continue
+                        new_event = ["key_signature", 0, 0, tr, (0 if (len(cs) == 1 and cs[0] == 9) else sf) + 7, 0]
+                        event_list.append(new_event)
+                else:
+                    for key_sig in key_sigs:
+                        tr = key_sig[3]
+                        if tr in track_to_channels:
+                            cs = track_to_channels[tr]
+                            if len(cs) == 1 and cs[0] == 9:
+                                continue
+                        key_sig[4] = sf + 7
+                        key_sig[5] = 0
+            else:
+                # remove default key signature
+                for key_sig in key_sigs:
+                    event_list.remove(key_sig)
 
         events_name_order = ["time_signature", "key_signature", "set_tempo", "patch_change", "control_change", "note"]
         events_name_order = {name: i for i, name in enumerate(events_name_order)}
@@ -830,7 +863,10 @@ class MIDITokenizerV2:
             else:
                 if event[0] == "note":
                     notes_in_setup = True
-                key = tuple(event[3:-1])
+                if event[0] in ["note", "time_signature", "key_signature"]:
+                    key = tuple([event[0]]+event[3:-2])
+                else:
+                    key = tuple([event[0]]+event[3:-1])
             setup_events[key] = new_event
 
         last_t1 = 0