1.reduce the #selected features. 2.better feature format. 3.add normalize_rag_feature arg when creating/loading rag index db

Author: nbbb24

ecg_bench/config.py

@@ -60,6 +60,7 @@ def get_args():
     mode_group.add_argument('--retrieved_information', type=str, default='combined', choices=['feature', 'report', 'combined'], help='Type of information to retrieve in output')
     mode_group.add_argument('--load_rag_db', type = str, default = None, help = 'Load a RAG database')
     mode_group.add_argument('--load_rag_db_idx', type = str, default = None, help = 'Load a RAG database index')
+    mode_group.add_argument('--normalized_rag_feature', action='store_true', default=True, help='Enable normalization for RAG features and signals')
     mode_group.add_argument('--dev', action='store_true', default=None, help='Development mode')
     mode_group.add_argument('--log', action='store_true', default=None, help='Enable logging')
 

ecg_bench/utils/data_loader_utils.py

@@ -110,7 +110,11 @@ def setup_conversation_template(self, signal = None):
         if self.args.retrieval_base in ['feature', 'combined']:
             if self.args.dev:
                 print("🔍 DEBUG: Extracting features")
-            feature=self.rag_db.feature_extractor.extract_features(signal)
+            original_feature=self.rag_db.feature_extractor.extract_rag_features(signal)
+            feature=original_feature
+            if self.args.normalized_rag_feature:
+                feature = self.rag_db.query_feature_normalization(original_feature)
+                signal = self.rag_db.query_signal_lead_normalization(signal)
             if self.args.dev:
                 print("🔍 DEBUG: Features extracted, shape: ", feature.shape)
 
@@ -161,13 +165,17 @@ def append_messages_to_conv(self, conv, altered_text, signal=None):
             message_value = message_value.replace('<ecg>', '')
             message_value = message_value.replace('image', 'signal').replace('Image', 'Signal')
             if self.args.retrieval_base in ['feature', 'combined'] or self.args.retrieved_information in ['feature','combined']:
-                feature=self.rag_db.feature_extractor.extract_features(signal)
+                original_feature=self.rag_db.feature_extractor.extract_rag_features(signal)
+                feature=original_feature
+                if self.args.normalized_rag_feature:
+                    feature = self.rag_db.query_feature_normalization(original_feature)
+                    signal = self.rag_db.query_signal_lead_normalization(signal)
             if is_human and count == 0:
                 if self.args.rag and self.args.rag_prompt_mode == 'user_query':
                     rag_results = self.rag_db.search_similar(query_features=feature, query_signal=signal, k=self.args.rag_k, mode=self.args.retrieval_base)
                     filtered_rag_results = self.rag_db.format_search(rag_results,self.args.retrieved_information)
                     if self.args.retrieved_information == 'combined':
-                        message_value = f"<signal>\nFeature Information:\n{feature}\n\n{filtered_rag_results}\n{message_value}"
+                        message_value = f"<signal>\nFeature Information:\n{self.rag_db.convert_features_to_structured(original_feature)}\n\n{filtered_rag_results}\n{message_value}"
                     elif self.args.retrieved_information == 'report':
                         message_value = f"<signal>\n{filtered_rag_results}\n{message_value}"
                 else:

ecg_bench/utils/preprocess_utils.py

@@ -1047,6 +1047,61 @@ def extract_features(self, ecg):
 
         return np.array(features)
 
+    def extract_rag_features(self, ecg):
+        """
+        Extract a subset of features for RAG applications.
+        Keeps only: max, min, dominant_frequency, total_power, spectral_centroid, 
+        peak_frequency_power, Heart Rate Features, Wavelet Features, average_absolute_difference, root_mean_square_difference
+        """
+        features = []
+        
+        for lead in range(ecg.shape[0]):
+            lead_signal = ecg[lead, :]
+            
+            # Basic statistical features (only max and min)
+            features.extend([
+                np.max(lead_signal),
+                np.min(lead_signal)
+            ])
+            
+            # Frequency domain features
+            freqs, psd = signal.welch(lead_signal, fs=self.target_sf, nperseg=min(1024, len(lead_signal)))
+            total_power = np.sum(psd)
+            features.extend([
+                total_power,  # Total power
+                np.max(psd),  # Peak frequency power
+                freqs[np.argmax(psd)],  # Dominant frequency
+            ])
+            
+            # Spectral centroid with NaN handling
+            if total_power > 0:
+                spectral_centroid = np.sum(freqs * psd) / total_power
+            else:
+                spectral_centroid = 0
+            features.append(spectral_centroid)
+            
+            # Find peaks with robust thresholding
+            if np.max(lead_signal) != np.min(lead_signal):  # Avoid division by zero
+                peak_height = 0.3 * (np.max(lead_signal) - np.min(lead_signal)) + np.min(lead_signal)
+                min_distance = max(int(0.2 * self.target_sf), 1)  # Ensure positive distance
+                peaks, _ = signal.find_peaks(lead_signal, height=peak_height, distance=min_distance)
+            else:
+                peaks = []
+                
+            # Heart rate features
+            heart_rate_features = self._calculate_heart_rate_features(lead_signal, peaks)
+            features.extend(heart_rate_features)
+            
+            # Wavelet features
+            wavelet_features = self._calculate_wavelet_features(lead_signal)
+            features.extend(wavelet_features)
+            
+            # Non-linear features
+            features.append(np.mean(np.abs(np.diff(lead_signal))))  # Average absolute difference
+            features.append(np.sqrt(np.mean(np.square(np.diff(lead_signal)))))  # Root mean square of successive differences
+        
+        return np.array(features)
+    
     def _calculate_heart_rate_features(self, ecg, peaks):
         if len(peaks) > 1:
             # Heart rate
@@ -1118,4 +1173,58 @@ def find_st_deviation(self, ecg, peaks):
             return ecg[st_point] - ecg[peaks[-1]]
         return 0
 
+    def signal_lead_normalization(ecg):
+        """
+        Normalize each lead individually using z-score normalization.
+        """
+        if ecg.shape[0] == 12: 
+            ecg = ecg.T
+            transpose_back = True
+        else:
+            transpose_back = False
+        
+        normalized_ecg = np.zeros_like(ecg, dtype=np.float32)
+        
+        for lead_idx in range(12):
+            lead_signal = ecg[:, lead_idx]
+            lead_mean = np.mean(lead_signal)
+            lead_std = np.std(lead_signal) + 1e-10
+            normalized_ecg[:, lead_idx] = (lead_signal - lead_mean) / lead_std
+
+        if transpose_back:
+            normalized_ecg = normalized_ecg.T
+        
+        return normalized_ecg
+        
+    def feature_normalization(self, rag_features):
+        """
+        Normalize RAG features using z-score normalization.
+        """
+        features_per_lead = len(self.ecg_feature_list)
+        expected_total_features = 12 * features_per_lead
+        
+        if rag_features.ndim != 1:
+            raise ValueError(f"Expected 1D array, got shape {rag_features.shape}")
+        
+        if len(rag_features) != expected_total_features:
+            raise ValueError(f"Expected {expected_total_features} features for 12-lead ECG, got {len(rag_features)}")
+        
+        normalized_features = np.zeros_like(rag_features, dtype=np.float32)
+        
+        for feature_idx, feature_name in enumerate(self.ecg_feature_list):
+            feature_values = []
+            for lead_idx in range(12):
+                feature_pos = lead_idx * features_per_lead + feature_idx
+                feature_values.append(rag_features[feature_pos])
+            
+            feature_values = np.array(feature_values)
+            
+            feature_mean = np.mean(feature_values)
+            feature_std = np.std(feature_values) + 1e-10 
+            
+            for lead_idx in range(12):
+                feature_pos = lead_idx * features_per_lead + feature_idx
+                normalized_features[feature_pos] = (rag_features[feature_pos] - feature_mean) / feature_std
+        
+        return normalized_features