Add support for sustain pedal (pianoroll and chromagram) (#142)

* Added support for elongating note when encountering CC64

*  - Add unit-tests

*  - Add pedal for chroma and piano roll (#130)

* revised #130 - documented the code and replaced pedal logic with a simpler algorithm

* Revise #130 - unify use of pedals to single parameter pedal_threshold

* Fix unit-tests for #130

* Fix pep8 formatting issues

* Set the default pedal value to 64

* Modified tests to match the current specifications

* Use >= pedal_threshold as a condition for detecting pedal-on

pretty_midi/instrument.py

@@ -73,7 +73,8 @@ def get_onsets(self):
         # Return them sorted (because why not?)
         return np.sort(onsets)
 
-    def get_piano_roll(self, fs=100, times=None):
+    def get_piano_roll(self, fs=100, times=None,
+                       pedal_threshold=64):
         """Compute a piano roll matrix of this instrument.
 
         Parameters
@@ -84,6 +85,12 @@ def get_piano_roll(self, fs=100, times=None):
         times : np.ndarray
             Times of the start of each column in the piano roll.
             Default ``None`` which is ``np.arange(0, get_end_time(), 1./fs)``.
+        pedal_threshold : int
+            Value of control change 64 (sustain pedal) message that is less
+            than this value is reflected as pedal-off.  Pedals will be
+            reflected as elongation of notes in the piano roll.
+            If None, then CC64 message is ignored.
+            Default is 64.
 
         Returns
         -------
@@ -113,6 +120,29 @@ def get_piano_roll(self, fs=100, times=None):
             piano_roll[note.pitch,
                        int(note.start*fs):int(note.end*fs)] += note.velocity
 
+        # Process sustain pedals
+        if pedal_threshold is not None:
+            CC_SUSTAIN_PEDAL = 64
+            time_pedal_on = 0
+            is_pedal_on = False
+            for cc in [_e for _e in self.control_changes
+                       if _e.number == CC_SUSTAIN_PEDAL]:
+                time_now = int(cc.time*fs)
+                is_current_pedal_on = (cc.value >= pedal_threshold)
+                if not is_pedal_on and is_current_pedal_on:
+                    time_pedal_on = time_now
+                    is_pedal_on = True
+                elif is_pedal_on and not is_current_pedal_on:
+                    # For each pitch, a sustain pedal "retains"
+                    # the maximum velocity up to now due to
+                    # logarithmic nature of human loudness perception
+                    subpr = piano_roll[:, time_pedal_on:time_now]
+
+                    # Take the running maximum
+                    pedaled = np.maximum.accumulate(subpr, axis=1)
+                    piano_roll[:, time_pedal_on:time_now] = pedaled
+                    is_pedal_on = False
+
         # Process pitch changes
         # Need to sort the pitch bend list for the following to work
         ordered_bends = sorted(self.pitch_bends, key=lambda bend: bend.time)
@@ -163,7 +193,7 @@ def get_piano_roll(self, fs=100, times=None):
                                                   axis=1)
         return piano_roll_integrated
 
-    def get_chroma(self, fs=100, times=None):
+    def get_chroma(self, fs=100, times=None, pedal_threshold=64):
         """Get a sequence of chroma vectors from this instrument.
 
         Parameters
@@ -174,6 +204,12 @@ def get_chroma(self, fs=100, times=None):
         times : np.ndarray
             Times of the start of each column in the piano roll.
             Default ``None`` which is ``np.arange(0, get_end_time(), 1./fs)``.
+        pedal_threshold : int
+            Value of control change 64 (sustain pedal) message that is less
+            than this value is reflected as pedal-off.  Pedals will be
+            reflected as elongation of notes in the piano roll.
+            If None, then CC64 message is ignored.
+            Default is 64.
 
         Returns
         -------
@@ -182,7 +218,8 @@ def get_chroma(self, fs=100, times=None):
 
         """
         # First, get the piano roll
-        piano_roll = self.get_piano_roll(fs=fs, times=times)
+        piano_roll = self.get_piano_roll(fs=fs, times=times,
+                                         pedal_threshold=pedal_threshold)
         # Fold into one octave
         chroma_matrix = np.zeros((12, piano_roll.shape[1]))
         for note in range(12):

pretty_midi/pretty_midi.py

@@ -739,7 +739,7 @@ def get_onsets(self):
         # Return them sorted (because why not?)
         return np.sort(onsets)
 
-    def get_piano_roll(self, fs=100, times=None):
+    def get_piano_roll(self, fs=100, times=None, pedal_threshold=64):
         """Compute a piano roll matrix of the MIDI data.
 
         Parameters
@@ -750,6 +750,12 @@ def get_piano_roll(self, fs=100, times=None):
         times : np.ndarray
             Times of the start of each column in the piano roll.
             Default ``None`` which is ``np.arange(0, get_end_time(), 1./fs)``.
+        pedal_threshold : int
+            Value of control change 64 (sustain pedal) message that is less
+            than this value is reflected as pedal-off.  Pedals will be
+            reflected as elongation of notes in the piano roll.
+            If None, then CC64 message is ignored.
+            Default is 64.
 
         Returns
         -------
@@ -763,7 +769,8 @@ def get_piano_roll(self, fs=100, times=None):
             return np.zeros((128, 0))
 
         # Get piano rolls for each instrument
-        piano_rolls = [i.get_piano_roll(fs=fs, times=times)
+        piano_rolls = [i.get_piano_roll(fs=fs, times=times,
+                                        pedal_threshold=pedal_threshold)
                        for i in self.instruments]
         # Allocate piano roll,
         # number of columns is max of # of columns in all piano rolls
@@ -833,7 +840,7 @@ def get_pitch_class_transition_matrix(self, normalize=False,
 
         return pc_trans_mat
 
-    def get_chroma(self, fs=100, times=None):
+    def get_chroma(self, fs=100, times=None, pedal_threshold=64):
         """Get the MIDI data as a sequence of chroma vectors.
 
         Parameters
@@ -844,6 +851,12 @@ def get_chroma(self, fs=100, times=None):
         times : np.ndarray
             Times of the start of each column in the piano roll.
             Default ``None`` which is ``np.arange(0, get_end_time(), 1./fs)``.
+        pedal_threshold : int
+            Value of control change 64 (sustain pedal) message that is less
+            than this value is reflected as pedal-off.  Pedals will be
+            reflected as elongation of notes in the piano roll.
+            If None, then CC64 message is ignored.
+            Default is 64.
 
         Returns
         -------
@@ -852,7 +865,8 @@ def get_chroma(self, fs=100, times=None):
 
         """
         # First, get the piano roll
-        piano_roll = self.get_piano_roll(fs=fs, times=times)
+        piano_roll = self.get_piano_roll(fs=fs, times=times,
+                                         pedal_threshold=pedal_threshold)
         # Fold into one octave
         chroma_matrix = np.zeros((12, piano_roll.shape[1]))
         for note in range(12):

tests/test_pretty_midi.py

@@ -406,19 +406,58 @@ def test_get_piano_roll_and_get_chroma():
     inst.notes.append(pretty_midi.Note(pitch=45, velocity=100, start=0.1,
                                        end=0.2))
 
+    inst = pretty_midi.Instrument(0)
+    pm.instruments.append(inst)
+    inst.control_changes.append(pretty_midi.ControlChange(number=64,
+                                                          value=65,
+                                                          time=0.12))
+    inst.control_changes.append(pretty_midi.ControlChange(number=64,
+                                                          value=63,
+                                                          time=0.5))
+    inst.notes.append(pretty_midi.Note(pitch=50, velocity=50, start=0.35,
+                                       end=0.4))
+    inst.notes.append(pretty_midi.Note(pitch=55, velocity=20, start=0.1,
+                                       end=0.15))
+    inst.notes.append(pretty_midi.Note(pitch=55, velocity=10, start=0.2,
+                                       end=0.25))
+    inst.notes.append(pretty_midi.Note(pitch=55, velocity=50, start=0.3,
+                                       end=0.42))
+
     expected_piano_roll = np.zeros((128, 50))
     expected_piano_roll[40, 5:35] = 100
     expected_piano_roll[40, 35:45] = 150
     expected_piano_roll[40, 45:] = 50
     expected_piano_roll[45, 10:20] = 100
-    assert np.allclose(pm.get_piano_roll(), expected_piano_roll)
+    expected_piano_roll[50, 35:40] = 50
+    expected_piano_roll[55, 10:15] = 20
+    expected_piano_roll[55, 20:25] = 10
+    expected_piano_roll[55, 30:42] = 50
+    assert np.allclose(pm.get_piano_roll(pedal_threshold=None),
+                       expected_piano_roll)
+
+    expected_piano_roll[50, 35:50] = 50
+    expected_piano_roll[55, 10:30] = 20
+    expected_piano_roll[55, 30:50] = 50
+    assert np.allclose(pm.get_piano_roll(),
+                       expected_piano_roll)
 
     expected_chroma = np.zeros((12, 50))
     expected_chroma[4, 5:35] = 100
     expected_chroma[4, 35:45] = 150
     expected_chroma[4, 45:] = 50
     expected_chroma[9, 10:20] = 100
-    assert np.allclose(pm.get_chroma(), expected_chroma)
+    expected_chroma[2, 35:40] = 50
+    expected_chroma[7, 10:15] = 20
+    expected_chroma[7, 20:25] = 10
+    expected_chroma[7, 30:42] = 50
+    assert np.allclose(pm.get_chroma(pedal_threshold=None),
+                       expected_chroma)
+
+    expected_chroma[2, 35:50] = 50
+    expected_chroma[7, 10:30] = 20
+    expected_chroma[7, 30:50] = 50
+    assert np.allclose(pm.get_chroma(),
+                       expected_chroma)
 
 
 def test_synthesize():