add shiftsamps option for rdann

cx1111 · cx1111 · commit ea5812f2b62d · 2017-06-20T16:41:52.000-04:00
diff --git a/README.rst b/README.rst
@@ -243,7 +243,7 @@ Reading Annotations
 
 ::
 
-    annotation = rdann(recordname, annotator, sampfrom=0, sampto=None, pbdir=None)
+    annotation = rdann(recordname, annotator, sampfrom=0, sampto=None, shiftsamps=False, pbdir=None)
 
 Example Usage:
 ::
@@ -258,6 +258,7 @@ Input arguments:
   file '100.atr', annotator='atr'
 - ``sampfrom`` (default=0): The minimum sample number for annotations to be returned.
 - ``sampto`` (default=None): The maximum sample number for annotations to be returned.
+- ``shiftsamps`` (default=False): Boolean flag that specifies whether to return the sample indices relative to 'sampfrom' (True), or sample 0 (False). Annotation files store exact sample locations. 
 - ``pbdir`` (default=None): Option used to stream data from Physiobank. The Physiobank database directory from which to find the required annotation file. eg. For record '100' in 'http://physionet.org/physiobank/database/mitdb', pbdir = 'mitdb'.
 
 Output arguments:
diff --git a/wfdb/readwrite/annotations.py b/wfdb/readwrite/annotations.py
@@ -587,7 +587,7 @@ def showanncodes():
 ## ------------- Reading Annotations ------------- ##
 
 
-def rdann(recordname, annotator, sampfrom=0, sampto=None, pbdir=None):
+def rdann(recordname, annotator, sampfrom=0, sampto=None, shiftsamps=False, pbdir=None):
     """ Read a WFDB annotation file recordname.annotator and return an
     Annotation object.
 
@@ -601,6 +601,9 @@ def rdann(recordname, annotator, sampfrom=0, sampto=None, pbdir=None):
       file '100.atr', annotator='atr'
     - sampfrom (default=0): The minimum sample number for annotations to be returned.
     - sampto (default=None): The maximum sample number for annotations to be returned.
+    - shiftsamps (default=False): Boolean flag that specifies whether to return the
+      sample indices relative to 'sampfrom' (True), or sample 0 (False). Annotation files
+      store exact sample locations.
     - pbdir (default=None): Option used to stream data from Physiobank. The Physiobank database 
       directory from which to find the required annotation file.
       eg. For record '100' in 'http://physionet.org/physiobank/database/mitdb', pbdir = 'mitdb'.
@@ -657,7 +660,6 @@ def rdann(recordname, annotator, sampfrom=0, sampto=None, pbdir=None):
     # Set the annotation type to the annotation codes
     anntype = [allannsyms[code] for code in anntype]
 
-
     # If the fs field was not present in the file, try to read a wfdb header
     # the annotation is associated with to get the fs
     if fs is None:
@@ -668,6 +670,11 @@ def rdann(recordname, annotator, sampfrom=0, sampto=None, pbdir=None):
         except:
             pass
 
+    # Return annotation samples relative to starting signal index
+    if shiftsamps and annsamp!=[] and sampfrom:
+        annsamp = np.array(annsamp, dtype='int64')
+        annsamp = annsamp - sampfrom
+
     # Store fields in an Annotation object
     annotation = Annotation(os.path.split(recordname)[1], annotator, annsamp, anntype, 
         subtype, chan, num, aux, fs,custom_anntypes)
