Update prediction.py

prediction.py

@@ -21,52 +21,60 @@ def aggregate_predictions(spectrogram_list, model, threshold=0.5):
     return final_label, mean_prob, segment_probs
 
 def predict_eeg_recording(edf_path, model_name='2DCNN', threshold=0.5):
-    print(f"\n=== [DEBUG] Loading model: {model_name} ===")
-
     if model_name == '2DCNN':
         model = tf.keras.models.load_model('model1_2dcnn.h5')
         channels = ["EEG FP1-REF", "EEG FP2-REF", "EEG F3-REF", "EEG F4-REF", "EEG C3-REF"]
-        
         preprocessed_df = preprocess_eeg_file(edf_path, fmin=1.0, fmax=45.0, segment_lenght=5, overlap=2, desired=channels)
-        print(f"[2DCNN] Nombre de segments : {len(preprocessed_df)}")
 
         if preprocessed_df is None or preprocessed_df.empty:
             raise ValueError("EEG file could not be preprocessed or no valid segments found.")
-        
+
         spectrogram_list = preprocessed_df.apply(
             lambda row: convert_epoch_to_spectrogram(row, channels, fs=250, nperseg=128, noverlap=64), axis=1
         ).tolist()
 
-        print(f"[2DCNN] Shape segment 0 : {spectrogram_list[0].shape}")
         return aggregate_predictions(spectrogram_list, model, threshold)
 
     elif model_name == 'EEGNet':
         loaded = joblib.load("eegnet_model.joblib")
         state_dict = loaded["model_state_dict"]
-        
+
         model = EEGNet(n_channels=21, n_samples=1250, num_classes=2)
         model.load_state_dict(state_dict)
-        
-        channels = [ ... ]  # ta liste de 21 canaux
-        preprocessed_df = preprocess_eeg_file(edf_path, fmin=1.0, fmax=45.0, segment_lenght=5, overlap=0, desired=channels)
-        print(f"[EEGNet] Nombre de segments : {len(preprocessed_df)}")
+
+        channels = [
+            'EEG FP1-REF', 'EEG FP2-REF', 'EEG F3-REF', 'EEG F4-REF', 'EEG C3-REF', 'EEG C4-REF',
+            'EEG P3-REF', 'EEG P4-REF', 'EEG O1-REF', 'EEG O2-REF', 'EEG F7-REF', 'EEG F8-REF',
+            'EEG T3-REF', 'EEG T4-REF', 'EEG T5-REF', 'EEG T6-REF', 'EEG FZ-REF', 'EEG CZ-REF',
+            'EEG PZ-REF', 'EEG ROC-REF', 'EEG LOC-REF'
+        ]
+
+        preprocessed_df = preprocess_eeg_file(
+            edf_path, fmin=1.0, fmax=45.0, segment_lenght=5, overlap=0, desired=channels
+        )
 
         if preprocessed_df is None or preprocessed_df.empty:
             raise ValueError("EEG file could not be preprocessed or no valid segments found.")
-        
+
         timeseries_list = preprocessed_df.apply(
             lambda row: convert_epoch_to_timeseries(row, channels), axis=1
         ).tolist()
 
-        print(f"[EEGNet] Shape segment 0 : {timeseries_list[0].shape}")
         return aggregate_predictions_EEGNET(timeseries_list, model, threshold)
 
     elif model_name == 'EpilepsyNet':
-        print("[EpilepsyNet] Chargement du fichier EDF brut avec MNE...")
         raw = mne.io.read_raw_edf(edf_path, preload=True, verbose='ERROR')
-
-        eeg_cols = [ ... ]  # ta liste de 21 colonnes
-        parameters = { ... }  # ton dictionnaire de paramètres
+        eeg_cols = ['EEG FP1', 'EEG FP2', 'EEG F3', 'EEG F4', 'EEG C3', 'EEG C4', 'EEG P3', 'EEG P4', 
+                    'EEG O1', 'EEG O2', 'EEG F7', 'EEG F8', 'EEG T3', 'EEG T4', 'EEG T5', 'EEG T6', 
+                    'EEG T1', 'EEG T2', 'EEG FZ', 'EEG CZ', 'EEG PZ']
+
+        parameters = {
+            'eeg_cols': eeg_cols,
+            'segment_duration': 60.0,
+            'n_segments_per_file': 12,
+            'samples_per_segment': 1250,
+            'random_state': 42
+        }
 
         X = process_raw_files(
             raw_file=raw,
@@ -76,25 +84,26 @@ def predict_eeg_recording(edf_path, model_name='2DCNN', threshold=0.5):
             random_state=parameters['random_state']
         )
 
-        print(f"[EpilepsyNet] Shape X : {X.shape}")
-
         X_std = standardize_data(X)
         corr_matrix = compute_correlation_matrix(X_std)
         upper_triangle_matrix = extract_upper_triangle(corr_matrix)
-        print(f"[EpilepsyNet] Shape upper_triangle_matrix : {upper_triangle_matrix.shape}")
 
-        X_tensor = torch.tensor(upper_triangle_matrix, dtype=torch.float32).unsqueeze(0)
+        X_tensor = torch.tensor(upper_triangle_matrix, dtype=torch.float32)
+        X_tensor = X_tensor.unsqueeze(0)
+
+        input_dim = 210
+        embed_dim = 256
+        num_heads = 16
 
-        model = TimeSeriesAttentionClassifier(input_dim=210, embed_dim=256, num_heads=16)
+        model = TimeSeriesAttentionClassifier(input_dim, embed_dim, num_heads)
         model.load_state_dict(torch.load('EpilepsyNet.pth'))
         model.eval()
 
         outputs, _ = model(X_tensor)
         predicted = (outputs >= 0.5).float()
-        
-        print(f"[EpilepsyNet] Probabilité brute : {outputs.squeeze().item()}")
 
-        return int(predicted), outputs.float().squeeze().item()
+        prob = outputs.float().squeeze().item()
+        return int(predicted), prob, [prob]
 
 def convert_epoch_to_timeseries(epoch_row, channels):
     ts_list = []