add quantize_notes function

CHANGES.rst

@@ -7,6 +7,7 @@ Version 0.16.dev0
 New features:
 
 * Added chord evaluation (#309)
+* Added `quantize_notes` function (#327)
 
 Bug fixes:
 

madmom/utils/__init__.py

@@ -388,9 +388,10 @@ def quantize_events(events, fps, length=None, shift=None):
     fps : float
         Quantize with `fps` frames per second.
     length : int, optional
-        Length of the returned array.
+        Length of the returned array. If 'None', the length will be set
+        according to the latest event.
     shift : float, optional
-        Shift the events by this value before quantization.
+        Shift the events by `shift` seconds before quantization.
 
     Returns
     -------
@@ -425,6 +426,93 @@ def quantize_events(events, fps, length=None, shift=None):
     return quantized
 
 
+def quantize_notes(notes, fps, length=None, num_notes=None, shift=None,
+                   offset=None):
+    """
+    Quantize the notes with the given resolution.
+
+    Parameters
+    ----------
+    notes : 2D numpy array
+        Notes to be quantized.
+    fps : float
+        Quantize with `fps` frames per second.
+    length : int, optional
+        Length of the returned array. If 'None', the length will be set
+        according to the latest sounding note.
+    num_notes : int, optional
+        Number of notes of the returned array. If 'None', the number of notes
+        will be based on the highest note number in the `notes` array. All note
+        numbers equal or greater this value will not be included in the
+        returned quantized array.
+    shift : float, optional
+        Shift the events by `shift` seconds before quantization.
+    offset : int or numpy array, optional
+        Apply this offset to the note numbers. If an integer value is given it
+        will be used for all note numbers. An array can be used to map notes
+        numbers arbitrarily. If a note number becomes negative, it will be not
+        included in the returned quantized array.
+
+    Returns
+    -------
+    numpy array
+        Quantized notes.
+
+    Notes
+    -----
+    Expected notes format:
+
+    'note_time' 'note_number' ['duration' ['velocity']]
+
+    """
+    # convert to numpy array or create a copy if needed
+    notes = np.array(notes, dtype=np.float, ndmin=2)
+    # check supported dims and shapes
+    if notes.ndim != 2:
+        raise ValueError('only 2-dimensional notes supported.')
+    if notes.shape[1] not in (2, 3, 4):
+        raise ValueError('2, 3, or 4 columns supported.')
+    # split the notes into columns
+    note_onsets = notes[:, 0]
+    note_numbers = notes[:, 1].astype(np.int)
+    # shift note onsets (before inferring offsets if needed)
+    if shift:
+        note_onsets += shift
+    note_offsets = np.copy(note_onsets)
+    if notes.shape[1] > 2:
+        note_offsets += notes[:, 2]
+    if notes.shape[1] > 3:
+        note_velocities = notes[:, 3]
+    else:
+        note_velocities = np.ones(len(notes))
+    # determine length of quantized array
+    if length is None:
+        # set the length to be long enough to cover all notes
+        length = int(round(np.max(note_offsets) * float(fps))) + 1
+    # offset note numbers
+    if offset is not None:
+        try:
+            note_numbers += offset[note_numbers]
+        except TypeError:
+            note_numbers += offset
+    # determine width of quantized array
+    if num_notes is None:
+        num_notes = int(np.max(note_numbers)) + 1
+    # init array
+    quantized = np.zeros((length, num_notes))
+    # quantize
+    note_onsets = np.round((note_onsets * fps)).astype(np.int)
+    note_offsets = np.round((note_offsets * fps)).astype(np.int) + 1
+    # iterate over all notes
+    for n, note in enumerate(notes):
+        # use only the notes which fit in the array and note number >= 0
+        if num_notes > note_numbers[n] >= 0:
+            quantized[note_onsets[n]:note_offsets[n], note_numbers[n]] = \
+                note_velocities[n]
+    # return quantized array
+    return quantized
+
+
 # argparse action to set and overwrite default lists
 class OverrideDefaultListAction(argparse.Action):
     """

tests/test_utils.py

@@ -13,6 +13,7 @@
 from madmom.utils import *
 from . import (DATA_PATH, AUDIO_PATH, ANNOTATIONS_PATH, ACTIVATIONS_PATH,
                DETECTIONS_PATH)
+from .test_features_notes import NOTES
 
 FILE_LIST = [pj(DATA_PATH, 'README'),
              pj(DATA_PATH, 'commented_txt'),
@@ -126,8 +127,6 @@ def test_recursion(self):
         result = search_path(DATA_PATH, 1)
         all_files = (FILE_LIST + AUDIO_FILES + ANNOTATION_FILES +
                      DETECTION_FILES + ACTIVATION_FILES)
-        print(len(result))
-        print(len(sorted(all_files)))
         self.assertEqual(result, sorted(all_files))
 
 
@@ -150,8 +149,6 @@ def test_file(self):
     def test_path(self):
         # no suffix
         result = search_files(DATA_PATH)
-        print("result", result)
-        print("FILELIST", sorted(FILE_LIST))
         self.assertEqual(result, sorted(FILE_LIST))
         # single suffix
         result = search_files(DATA_PATH, suffix='txt')
@@ -387,6 +384,104 @@ def test_shift(self):
         self.assertTrue(np.allclose(idx, correct))
 
 
+class TestQuantizeNotesFunction(unittest.TestCase):
+
+    def test_fps(self):
+        # 10 FPS
+        fps = 10
+        quantized = quantize_notes(NOTES, fps=fps)
+        self.assertTrue(quantized.shape == (42, 78))
+        idx = np.nonzero(quantized)
+        correct = np.arange(np.round(NOTES[0, 0] * fps),
+                            np.round((NOTES[0, 0] + NOTES[0, 2]) * fps) + 1)
+        self.assertTrue(np.allclose(idx[0][idx[1] == 72], correct))
+        # 100 FPS with numpy arrays (array must not be changed)
+        fps = 100
+        notes = np.array(NOTES)
+        notes_ = np.copy(notes)
+        quantized = quantize_notes(notes, fps=fps)
+        self.assertTrue(quantized.shape == (416, 78))
+        idx = np.nonzero(quantized)
+        correct = np.arange(np.round(NOTES[1, 0] * fps),
+                            np.round((NOTES[1, 0] + NOTES[1, 2]) * fps) + 1)
+        self.assertTrue(np.allclose(idx[0][idx[1] == 41], correct))
+        self.assertTrue(np.allclose(notes, notes_))
+
+    def test_length(self):
+        fps = 100
+        length = 280
+        quantized = quantize_notes(NOTES, fps=fps, length=length)
+        self.assertTrue(quantized.shape == (280, 78))
+        idx = np.nonzero(quantized)
+        correct = np.arange(np.round(NOTES[0, 0] * fps), length)
+        self.assertTrue(np.allclose(idx[0][idx[1] == 72], correct))
+
+    def test_rounding(self):
+        # rounding towards next even number
+        quantized = quantize_notes([[0.95, 0], [1.95, 1]], fps=10)
+        self.assertTrue(np.allclose(np.nonzero(quantized), [[10, 20], [0, 1]]))
+        quantized = quantize_notes([[0.85, 0], [1.85, 1]], fps=10)
+        self.assertTrue(np.allclose(np.nonzero(quantized), [[8, 18], [0, 1]]))
+        # round down
+        quantized = quantize_notes([[0.9499999, 0]], fps=10)
+        self.assertTrue(np.allclose(np.nonzero(quantized), [[9], [0]]))
+        # with length
+        quantized = quantize_notes([[0.95, 0], [1.95, 1]], fps=10, length=15)
+        self.assertTrue(np.allclose(np.nonzero(quantized), [[10], [0]]))
+
+    def test_shift(self):
+        fps = 10
+        shift = 1
+        quantized = quantize_notes(NOTES, fps=fps, shift=shift)
+        idx = np.nonzero(quantized)
+        correct = np.arange(np.round(NOTES[0, 0] * fps),
+                            np.round((NOTES[0, 0] + NOTES[0, 2]) * fps) + 1)
+        correct += shift * fps
+        self.assertTrue(np.allclose(idx[0][idx[1] == 72], correct))
+        # limited length
+        length = 35
+        quantized = quantize_notes(NOTES, fps=fps, shift=1, length=length)
+        idx = np.nonzero(quantized)
+        correct = np.arange(np.round(NOTES[0, 0] * fps) + shift * fps, length)
+        self.assertTrue(np.allclose(idx[0][idx[1] == 72], correct))
+
+    def test_offset(self):
+        fps = 10
+        offset = -21
+        quantized = quantize_notes(NOTES, fps=fps, offset=offset)
+        self.assertTrue(quantized.shape == (42, 57))
+        idx = np.nonzero(quantized)
+        correct = np.arange(np.round(NOTES[0, 0] * fps),
+                            np.round((NOTES[0, 0] + NOTES[0, 2]) * fps) + 1)
+        self.assertTrue(np.allclose(idx[0][idx[1] == 72 + offset], correct))
+        # negative note numbers will be ommited
+        offset = -60
+        quantized = quantize_notes(NOTES, fps=fps, offset=offset)
+        self.assertTrue(quantized.shape == (42, 18))
+        idx = np.nonzero(quantized)
+        correct = np.arange(np.round(NOTES[0, 0] * fps),
+                            np.round((NOTES[0, 0] + NOTES[0, 2]) * fps) + 1)
+        self.assertTrue(np.allclose(idx[0][idx[1] == 72 + offset], correct))
+
+    def test_num_notes(self):
+        fps = 10
+        # defining num_notes does only include notes up to this number (-1)
+        num_notes = 73
+        quantized = quantize_notes(NOTES, fps=fps, num_notes=num_notes)
+        self.assertTrue(quantized.shape == (42, 73))
+        idx = np.nonzero(quantized)
+        correct = np.arange(np.round(NOTES[0, 0] * fps),
+                            np.round((NOTES[0, 0] + NOTES[0, 2]) * fps) + 1)
+        self.assertTrue(np.allclose(idx[0][idx[1] == 72], correct))
+        # if we add an offset, note #72 will not be included any more
+        num_notes = 72
+        offset = 1
+        quantized = quantize_notes(NOTES, fps=fps, num_notes=num_notes,
+                                   offset=offset)
+        idx = np.nonzero(quantized)
+        self.assertTrue(np.allclose(idx[0][idx[1] == 72], []))
+
+
 class TestSegmentAxisFunction(unittest.TestCase):
 
     def test_types(self):