work on cqrs

cx1111 · cx1111 · commit f50e46c57764 · 2018-01-12T22:19:52.000-05:00
diff --git a/wfdb/processing/cqrs.py b/wfdb/processing/cqrs.py
@@ -11,8 +11,23 @@ class Conf(object):
     """
     Initial configuration object
     """
-    def __init__(self, hr=75, qrs_width=0.1):
+    def __init__(self, hr=75, hr_max=200, hr_min=25, qrs_width=0.1):
+        """
+        Parameters
+        ----------
+        hr : int or float, optional
+            Heart rate in beats per minute
+        qrs_width : int or float, optional
+            Expected qrs width in seconds
+        hr_max : int or float, optional
+            Hard maximum heart rate between two beats, in beats per minute
+        hr_min : int or float, optional
+            Hard minimum heart rate between two beats, in beats per minute
+        """
+
         self.hr = hr
+        self.hr_max = hr_max
+        self.hr_min = hr_min
         self.qrs_width = qrs_width
 
 
@@ -26,7 +41,10 @@ def __init__(self, sig, fs, conf):
         self.sig_len = len(sig)
 
         # These values are in samples
-        self.rr = 60 * fs / conf.hr
+        self.rr = 60 * self.fs / conf.hr
+        self.rr_max = 60 * self.fs / conf.hr_min
+        self.rr_min = 60 * self.fs / conf.hr_max
+
         self.qrs_width = conf.qrs_width * fs
 
 
@@ -54,24 +72,28 @@ def mwi(self):
     def learn_params(self):
         """
         Learn the following:
-        - rr interval
-        - 
+        - qrs threshold
 
         """
-        # Find the dominant peaks of the signal
-        dominant_peaks = self.sig_i[find_dominant_peaks(self.sig_i,
-                                             int(self.qrs_width / 2))]
-        # Cluster and use the largest cluster mean as expected param
-        kmeans = KMeans(n_clusters=2, random_state=0).fit(dominant_peaks)
+        # Find the dominant peaks of the signal. Store indices.
+        self.dominant_peak_inds = find_dominant_peaks(self.sig_i,
+                                                  int(self.qrs_width / 2))
+        dominant_peaks = self.sig_i[self.dominant_peak_inds]
 
-        # The mean value of the top cluster
+        # Cluster the peaks. The peaks in the largest amplitude cluster are
+        # treated as qrs peaks. Use these to determine qrs detection thresholds
+        kmeans = KMeans(n_clusters=2, random_state=0).fit(dominant_peaks)
         top_cluster_ind = np.argmax(kmeans.cluster_centers_)
+        
+        self.qrs_peak_inds = self.dominant_peak_inds[np.where(kmeans.labels_ == top_cluster_ind)]
+        qrs_peaks = self.sig_i[self.qrs_peak_inds]
 
-        self.qrs_thresh = max(kmeans.cluster_centers_)
-        # The standard deviation of the top cluster
-        self.qrs_dev = np.std(dominant_peaks[np.where(kmeans.labels_ == top_cluster_ind)])
+        self.qrs_thresh = np.mean(qrs_peaks) / 3
+        self.qrs_thresh_min = self.qrs_thresh / 3
 
-        pdb.set_trace()
+        return
+        
+        
 
     def detect(self, sampfrom=0, sampto='end'):
         """
@@ -88,6 +110,18 @@ def detect(self, sampfrom=0, sampto='end'):
 
         self.learn_params()
 
+
+        # plt.plot(self.sig_f, 'b')
+        # plt.plot(self.sig_i, 'r')
+        # #plt.plot(self.dominant_peak_inds, self.sig_i[self.dominant_peak_inds], 'k*')
+        # plt.plot(self.qrs_peak_inds, self.sig_i[self.qrs_peak_inds], 'g*')
+        # plt.title('Derived ECG signals')
+        # plt.legend(['Badpass filtered','Moving wave integration'])
+        # plt.show()
+
+        plt.plot(self.sig, 'b')
+        plt.plot(self.qrs_peak_inds, self.sig[self.qrs_peak_inds], 'r*')
+
         return
 
         for i in range(self.sig_len):
@@ -144,8 +178,6 @@ def cqrs_detect(sig, fs, conf=Conf()):
 
     cqrs.detect()
 
-    plt.plot(cqrs.sig_f, 'b')
-    plt.plot(cqrs.sig_i, 'r')
-    plt.show()
 
-    #return cqrs.annotation
+
+    return cqrs.qrs_inds
