Merge pull request #39 from bmcfee/pump

added Pump collector object

README.md

@@ -41,6 +41,11 @@ be easily consumed by statistical algorithms.  Some desired features:
 >>> p_beat = pumpp.task.BeatTransformer(sr=sr, hop_length=hop_length)
 >>> p_chord = pumpp.task.SimpleChordTransformer(sr=sr, hop_length=hop_length)
 
->>> # Apply the extractors
->>> data = pumpp.apply(audio_f, jams_f, p_cqt, p_beat, b_chord)
+>>> # Collect the operators
+>>> P = pumpp.Pump(p_cqt, p_beat, p_chord)
+>>> # Apply the extractors to generate training data
+>>> data = P.transform(audio_f, jams_f)
+
+>>> # Or test data
+>>> test_data = P.transform('/my/test/audio.ogg')
 ```

pumpp/core.py

@@ -6,12 +6,14 @@
 .. autosummary::
     :toctree: generated/
 
+    Pump
     transform
 '''
 
 import librosa
 import jams
 
+from .exceptions import ParameterError
 from .task import BaseTaskTransformer
 from .feature import FeatureExtractor
 
@@ -57,3 +59,75 @@ def transform(audio_f, jam, *ops):
         elif isinstance(op, FeatureExtractor):
             data.update(op.transform(y, sr))
     return data
+
+
+class Pump(object):
+    '''Top-level pump object.
+
+    This class is used to collect feature and task transformers
+
+    Attributes
+    ----------
+    ops : list of (BaseTaskTransformer, FeatureExtractor)
+        The operations to apply
+
+    Examples
+    --------
+    Create a CQT and chord transformer
+    >>> p_cqt = pumpp.feature.CQT('cqt', sr=44100, hop_length=1024)
+    >>> p_chord = pumpp.task.ChordTagTransformer(sr=44100, hop_length=1024)
+    >>> pump = pumpp.Pump(p_cqt, p_chord)
+    >>> data = pump.transform('/my/audio/file.mp3', '/my/jams/annotation.jams')
+
+
+    See Also
+    --------
+    transform
+    '''
+
+    def __init__(self, *ops):
+
+        self.ops = []
+        for op in ops:
+            self.add(op)
+
+    def add(self, op):
+        '''Add an operation to this pump.
+
+        Parameters
+        ----------
+        op : BaseTaskTransformer, FeatureExtractor
+            The operation to add
+
+        Raises
+        ------
+        ParameterError
+            if `op` is not of a correct type
+        '''
+        if not isinstance(op, (BaseTaskTransformer, FeatureExtractor)):
+            raise ParameterError('op={} must be one of '
+                                 '(BaseTaskTransformer, FeatureExtractor)'
+                                 .format(op))
+
+        self.ops.append(op)
+
+    def transform(self, audio_f, jam=None):
+        '''Apply the transformations to an audio file, and optionally JAMS object.
+
+        Parameters
+        ----------
+        audio_f : str
+            Path to audio file
+
+        jam : optional, `jams.JAMS`, str or file-like
+            Optional JAMS object/path to JAMS file/open file descriptor.
+
+            If provided, this will provide data for task transformers.
+
+        Returns
+        -------
+        data : dict
+            Data dictionary containing the transformed audio (and annotations)
+        '''
+
+        return transform(audio_f, jam, *self.ops)

tests/test_core.py

@@ -64,3 +64,69 @@ def test_transform(audio_f, jam, sr, hop_length):
     assert (np.abs(data['stft/mag'].shape[1] - data['beat/beat'].shape[1])
             * hop_length / float(sr)) <= 0.05
     pass
+
+
+@pytest.mark.parametrize('audio_f', ['tests/data/test.ogg'])
+def test_pump(audio_f, jam, sr, hop_length):
+
+    ops = [pumpp.feature.STFT(name='stft', sr=sr,
+                              hop_length=hop_length,
+                              n_fft=2*hop_length),
+
+           pumpp.task.BeatTransformer(name='beat', sr=sr,
+                                      hop_length=hop_length),
+
+           pumpp.task.ChordTransformer(name='chord', sr=sr,
+                                       hop_length=hop_length),
+
+           pumpp.task.StaticLabelTransformer(name='tags',
+                                             namespace='tag_open',
+                                             labels=['rock', 'jazz'])]
+
+    data1 = pumpp.transform(audio_f, jam, *ops)
+
+    pump = pumpp.Pump(*ops)
+    data2 = pump.transform(audio_f, jam)
+
+    assert data1.keys() == data2.keys()
+
+    for key in data1:
+        assert np.allclose(data1[key], data2[key])
+
+
+@pytest.mark.parametrize('audio_f', ['tests/data/test.ogg'])
+def test_pump_empty(audio_f, jam, sr, hop_length):
+
+    pump = pumpp.Pump()
+    data = pump.transform(audio_f, jam)
+    assert data == dict()
+
+
+def test_pump_add(sr, hop_length):
+
+    ops = [pumpp.feature.STFT(name='stft', sr=sr,
+                              hop_length=hop_length,
+                              n_fft=2*hop_length),
+
+           pumpp.task.BeatTransformer(name='beat', sr=sr,
+                                      hop_length=hop_length),
+
+           pumpp.task.ChordTransformer(name='chord', sr=sr,
+                                       hop_length=hop_length),
+
+           pumpp.task.StaticLabelTransformer(name='tags',
+                                             namespace='tag_open',
+                                             labels=['rock', 'jazz'])]
+
+    pump = pumpp.Pump()
+    assert pump.ops == []
+
+    for op in ops:
+        pump.add(op)
+        assert op in pump.ops
+
+
+@pytest.mark.xfail(raises=pumpp.ParameterError)
+def test_pump_add_bad():
+
+    pumpp.Pump('foo')