update clarity of benchmark_mitdb

cx1111 · cx1111 · commit baaa319e3ca8 · 2018-06-22T15:11:50.000-04:00
diff --git a/wfdb/processing/evaluate.py b/wfdb/processing/evaluate.py
@@ -380,7 +380,7 @@ def compare_annotations(ref_sample, test_sample, window_width, signal=None):
     return comparitor
 
 
-def benchmark_mitdb(detector, verbose=False):
+def benchmark_mitdb(detector, verbose=False, print_results=False):
     """
     Benchmark a qrs detector against mitdb's records.
 
@@ -390,11 +390,15 @@ def benchmark_mitdb(detector, verbose=False):
         The detector function.
     verbose : bool, optional
         The verbose option of the detector function.
+    print_results : bool, optional
+        Whether to print the overall performance, and the results for
+        each record.
 
     Returns
     -------
-    comparitors : list
-        List of Comparitor objects run on the records.
+    comparitors : dictionary
+        Dictionary of Comparitor objects run on the records, keyed on
+        the record names.
     specificity : float
         Aggregate specificity.
     positive_predictivity : float
@@ -434,8 +438,18 @@ def benchmark_mitdb(detector, verbose=False):
     false_positive_rate = np.mean(
         [c.false_positive_rate for c in comparitors])
 
+    comparitors = dict(zip(record_list, comparitors))
+
     print('Benchmark complete')
 
+    if print_results:
+        print('\nOverall MITDB Performance - Specificity: %.4f, Positive Predictivity: %.4f, False Positive Rate: %.4f\n'
+              % (specificity, positive_predictivity, false_positive_rate))
+        for record_name in record_list:
+            print('Record %s:' % record_name)
+            comparitors[record_name].print_summary()
+            print('\n\n')
+
     return comparitors, specificity, positive_predictivity, false_positive_rate
 
 
@@ -451,5 +465,6 @@ def benchmark_mitdb_record(rec, detector, verbose):
     comparitor = compare_annotations(ref_sample=ann_ref.sample[1:],
                                      test_sample=qrs_inds,
                                      window_width=int(0.1 * fields['fs']))
-    print('Finished record %s' % rec)
+    if verbose:
+        print('Finished record %s' % rec)
     return comparitor
diff --git a/wfdb/processing/qrs.py b/wfdb/processing/qrs.py
@@ -1,6 +1,7 @@
 import copy
-import numpy as np
+import pdb
 
+import numpy as np
 from scipy import signal
 from sklearn.preprocessing import normalize
 
@@ -146,7 +147,8 @@ def _set_conf(self):
 
     def _bandpass(self, fc_low=5, fc_high=20):
         """
-        Apply a bandpass filter onto the signal, and save the filtered signal.
+        Apply a bandpass filter onto the signal, and save the filtered
+        signal.
         """
         self.fc_low = fc_low
         self.fc_high = fc_high
@@ -162,24 +164,28 @@ def _bandpass(self, fc_low=5, fc_high=20):
 
     def _mwi(self):
         """
-        Apply moving wave integration with a ricker (Mexican hat) wavelet onto
-        the filtered signal, and save the square of the integrated signal.
+        Apply moving wave integration (mwi) with a ricker (Mexican hat)
+        wavelet onto the filtered signal, and save the square of the
+        integrated signal.
 
         The width of the hat is equal to the qrs width
 
-        Also find all local peaks in the mwi signal.
+        After integration, find all local peaks in the mwi signal.
         """
-        b = signal.ricker(self.qrs_width, 4)
-        self.sig_i = signal.filtfilt(b, [1], self.sig_f, axis=0) ** 2
+        wavelet_filter = signal.ricker(self.qrs_width, 8)
 
-        # Save the mwi gain  (x2 due to double filtering) and the total gain
-        # from raw to mwi
-        self.mwi_gain = get_filter_gain(b, [1],
+        self.sig_i = signal.filtfilt(wavelet_filter, [1], self.sig_f,
+                                     axis=0) ** 2
+
+        # Save the mwi gain (x2 due to double filtering) and the total
+        # gain from raw to mwi
+        self.mwi_gain = get_filter_gain(wavelet_filter, [1],
                          np.mean([self.fc_low, self.fc_high]), self.fs) * 2
         self.transform_gain = self.filter_gain * self.mwi_gain
         self.peak_inds_i = find_local_peaks(self.sig_i, radius=self.qrs_radius)
         self.n_peaks_i = len(self.peak_inds_i)
 
+
     def _learn_init_params(self, n_calib_beats=8):
         """
         Find a number of consecutive beats and use them to initialize:
