Add Resampling and normalization.

Dubrzr · Dubrzr · commit f7fe48ab872d · 2017-06-26T15:56:29.000Z
diff --git a/setup.py b/setup.py
@@ -57,6 +57,7 @@
         'matplotlib>=1.5.1',
         'requests>=2.10.0',
         'pandas>=0.19.1',
+        'scipy==0.19.0'
         ],
 
     # List additional groups of dependencies here (e.g. development
diff --git a/tests/test_processing.py b/tests/test_processing.py
@@ -0,0 +1,41 @@
+import wfdb
+import numpy
+
+class test_processing():
+
+    def test_1(self):
+        sig, fields = wfdb.srdsamp('sampledata/100')
+        ann = wfdb.rdann('sampledata/100', 'atr')
+
+        fs = fields['fs']
+        fs_target = 50
+
+        new_sig, new_ann = wfdb.processing.resample_singlechan(sig[:, 0], ann, fs, fs_target)
+
+        expected_length = int(sig.shape[0]*fs_target/fs)
+
+        assert new_sig.shape[0] == expected_length
+
+    def test_2(self):
+        sig, fields = wfdb.srdsamp('sampledata/100')
+        ann = wfdb.rdann('sampledata/100', 'atr')
+
+        fs = fields['fs']
+        fs_target = 50
+
+        new_sig, new_ann = wfdb.processing.resample_multichan(sig, ann, fs, fs_target)
+
+        expected_length = int(sig.shape[0]*fs_target/fs)
+
+        assert new_sig.shape[0] == expected_length
+        assert new_sig.shape[1] == sig.shape[1]
+
+    def test_3(self):
+        sig, _ = wfdb.srdsamp('sampledata/100')
+        lb = -5
+        ub = 15
+
+        x = wfdb.processing.normalize(sig[:, 0], lb, ub)
+        assert x.shape[0] == sig.shape[0]
+        assert numpy.min(x) >= lb
+        assert numpy.max(x) <= ub
diff --git a/wfdb/__init__.py b/wfdb/__init__.py
@@ -2,4 +2,5 @@
 from .readwrite._signals import estres, wrdatfile
 from .readwrite.annotations import Annotation, rdann, wrann, showanncodes
 from .readwrite.downloads import getdblist
-from .plot.plots import plotrec, plotann
+from .plot.plots import plotrec, plotann
+from . import processing
diff --git a/wfdb/processing/__init__.py b/wfdb/processing/__init__.py
@@ -0,0 +1 @@
+from .basic import resample_ann, resample_sig, resample_singlechan, resample_multichan, normalize
diff --git a/wfdb/processing/basic.py b/wfdb/processing/basic.py
@@ -0,0 +1,109 @@
+import numpy
+from scipy import signal
+
+from wfdb import Annotation
+
+
+def resample_ann(tt, annsamp):
+    # tt: numpy.array as returned by signal.resample
+    # annsamp: numpy.array containing indexes of annotations (Annotation.annsamp)
+
+    # Compute the new annotation indexes
+
+    result = numpy.zeros(len(tt), dtype='bool')
+    j = 0
+    tprec = tt[j]
+    for i, v in enumerate(annsamp):
+        while True:
+            d = False
+            if j+1 == len(tt):
+                result[j] = 1
+                break
+            tnow = tt[j+1]
+            if tprec <= v and v <= tnow:
+                if v-tprec < tnow-v:
+                    result[j] = 1
+                else:
+                    result[j+1] = 1
+                d = True
+            j += 1
+            tprec = tnow
+            if d:
+                break
+    return numpy.where(result==1)[0].astype('int64')
+
+
+def resample_sig(x, fs, fs_target):
+    # x: a numpy.array containing the signal
+    # fs: the current frequency
+    # fs_target: the target frequency
+
+    # Resample a signal
+
+    t = numpy.arange(x.shape[0]).astype('float64')
+
+    if fs == fs_target:
+        return x, t
+
+    new_length = int(x.shape[0]*fs_target/fs)
+    xx, tt = signal.resample(x, num=new_length, t=t)
+    assert xx.shape == tt.shape and xx.shape[0] == new_length
+    assert numpy.all(numpy.diff(tt) > 0)
+    return xx, tt
+
+
+def resample_singlechan(x, ann, fs, fs_target):
+    # x: a numpy.array containing the signal
+    # ann: an Annotation object
+    # fs: the current frequency
+    # fs_target: the target frequency
+
+    # Resample a single-channel signal with its annotations
+
+    xx, tt = resample_sig(x, fs, fs_target)
+
+    new_annsamp = resample_ann(tt, ann.annsamp)
+    assert ann.annsamp.shape == new_annsamp.shape
+
+    new_ann = Annotation(ann.recordname, ann.annotator, new_annsamp, ann.anntype, ann.num, ann.subtype, ann.chan, ann.aux, ann.fs)
+    return xx, new_ann
+
+
+def resample_multichan(xs, ann, fs, fs_target, resamp_ann_chan=0):
+    # xs: a numpy.ndarray containing the signals as returned by wfdb.srdsamp
+    # ann: an Annotation object
+    # fs: the current frequency
+    # fs_target: the target frequency
+    # resample_ann_channel: the signal channel that is used to compute new annotation indexes
+
+    # Resample multiple channels with their annotations
+
+    assert resamp_ann_chan < xs.shape[1]
+
+    lx = []
+    lt = None
+    for chan in range(xs.shape[1]):
+        xx, tt = resample_sig(xs[:, chan], fs, fs_target)
+        lx.append(xx)
+        if chan == resamp_ann_chan:
+            lt = tt
+
+    new_annsamp = resample_ann(lt, ann.annsamp)
+    assert ann.annsamp.shape == new_annsamp.shape
+
+    new_ann = Annotation(ann.recordname, ann.annotator, new_annsamp, ann.anntype, ann.num, ann.subtype, ann.chan, ann.aux, ann.fs)
+    return numpy.column_stack(lx), new_ann
+
+
+def normalize(x, lb=0, ub=1):
+    # lb: Lower bound
+    # ub: Upper bound
+
+    # Resizes a signal between the lower and upper bound
+
+    mid = ub - (ub - lb) / 2
+    min_v = numpy.min(x)
+    max_v = numpy.max(x)
+    mid_v =  max_v - (max_v - min_v) / 2
+    coef = (ub - lb) / (max_v - min_v)
+    return x * coef - (mid_v * coef) + mid

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+from .basic import resample_ann, resample_sig, resample_singlechan, resample_multichan, normalize`