Upload epilepsy_detection_model.py

v1/epilepsy_detection_model.py

@@ -0,0 +1,406 @@
+import numpy as np
+import pandas as pd
+import ast
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from sklearn.metrics import classification_report, confusion_matrix, roc_auc_score, f1_score
+import matplotlib.pyplot as plt
+import seaborn as sns
+from tqdm import tqdm
+import warnings
+warnings.filterwarnings('ignore')
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+print(f"Using device: {device}")
+
+class EpilepsyDataset(Dataset):
+    def __init__(self, csv_path):
+        self.data = pd.read_csv(csv_path)
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        # Parse the data string to list
+        data_list = ast.literal_eval(self.data.iloc[idx]['data'])
+        data_tensor = torch.FloatTensor(data_list)
+        label = torch.LongTensor([self.data.iloc[idx]['label']])[0]
+        return data_tensor, label
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, d_model, num_heads, dropout=0.1):
+        super().__init__()
+        assert d_model % num_heads == 0
+
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.d_k = d_model // num_heads
+
+        self.W_q = nn.Linear(d_model, d_model)
+        self.W_k = nn.Linear(d_model, d_model)
+        self.W_v = nn.Linear(d_model, d_model)
+        self.W_o = nn.Linear(d_model, d_model)
+
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        batch_size = x.size(0)
+
+        Q = self.W_q(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
+        K = self.W_k(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
+        V = self.W_v(x).view(batch_size, -1, self.num_heads, self.d_k).transpose(1, 2)
+
+        scores = torch.matmul(Q, K.transpose(-2, -1)) / np.sqrt(self.d_k)
+        attn_weights = F.softmax(scores, dim=-1)
+        attn_weights = self.dropout(attn_weights)
+
+        context = torch.matmul(attn_weights, V)
+        context = context.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)
+
+        output = self.W_o(context)
+        return output, attn_weights
+
+class TemporalConvBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, dilation, dropout=0.2):
+        super().__init__()
+        padding = (kernel_size - 1) * dilation // 2
+
+        self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size,
+                               padding=padding, dilation=dilation)
+        self.bn1 = nn.BatchNorm1d(out_channels)
+        self.relu1 = nn.ReLU()
+        self.dropout1 = nn.Dropout(dropout)
+
+        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size,
+                               padding=padding, dilation=dilation)
+        self.bn2 = nn.BatchNorm1d(out_channels)
+        self.relu2 = nn.ReLU()
+        self.dropout2 = nn.Dropout(dropout)
+
+        self.downsample = nn.Conv1d(in_channels, out_channels, 1) if in_channels != out_channels else None
+
+    def forward(self, x):
+        residual = x if self.downsample is None else self.downsample(x)
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu1(out)
+        out = self.dropout1(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        out = out + residual
+        out = self.relu2(out)
+        out = self.dropout2(out)
+
+        return out
+
+class ChannelAttention(nn.Module):
+    def __init__(self, channels, reduction=8):
+        super().__init__()
+        self.avg_pool = nn.AdaptiveAvgPool1d(1)
+        self.max_pool = nn.AdaptiveMaxPool1d(1)
+
+        self.fc = nn.Sequential(
+            nn.Linear(channels, channels // reduction, bias=False),
+            nn.ReLU(),
+            nn.Linear(channels // reduction, channels, bias=False),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        b, c, _ = x.size()
+
+        avg_out = self.fc(self.avg_pool(x).view(b, c))
+        max_out = self.fc(self.max_pool(x).view(b, c))
+
+        out = avg_out + max_out
+        return x * out.view(b, c, 1)
+
+class AdvancedEpilepsyDetector(nn.Module):
+    def __init__(self, input_dim=178, num_classes=2, dropout=0.3):
+        super().__init__()
+
+        self.input_proj = nn.Linear(input_dim, 256)
+        self.input_bn = nn.BatchNorm1d(256)
+
+        self.tcn_blocks = nn.ModuleList([
+            TemporalConvBlock(1, 64, kernel_size=7, dilation=1, dropout=dropout),
+            TemporalConvBlock(64, 128, kernel_size=5, dilation=2, dropout=dropout),
+            TemporalConvBlock(128, 256, kernel_size=3, dilation=4, dropout=dropout),
+            TemporalConvBlock(256, 256, kernel_size=3, dilation=8, dropout=dropout),
+        ])
+
+        self.channel_attn = ChannelAttention(256)
+
+        self.mha1 = MultiHeadAttention(256, num_heads=8, dropout=dropout)
+        self.mha2 = MultiHeadAttention(256, num_heads=8, dropout=dropout)
+
+        self.layer_norm1 = nn.LayerNorm(256)
+        self.layer_norm2 = nn.LayerNorm(256)
+
+        self.ffn = nn.Sequential(
+            nn.Linear(256, 512),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(512, 256),
+            nn.Dropout(dropout)
+        )
+
+        self.bilstm = nn.LSTM(256, 128, num_layers=2, batch_first=True,
+                              bidirectional=True, dropout=dropout)
+
+        self.classifier = nn.Sequential(
+            nn.Linear(256 + 256, 512),  # TCN output + LSTM output
+            nn.BatchNorm1d(512),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+
+            nn.Linear(512, 256),
+            nn.BatchNorm1d(256),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+
+            nn.Linear(256, 128),
+            nn.BatchNorm1d(128),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+
+            nn.Linear(128, num_classes)
+        )
+
+    def forward(self, x):
+        batch_size = x.size(0)
+
+        x_proj = self.input_proj(x)
+        x_proj = self.input_bn(x_proj)
+        x_proj = F.relu(x_proj)
+
+        x_tcn = x_proj.unsqueeze(1)
+        for tcn_block in self.tcn_blocks:
+            x_tcn = tcn_block(x_tcn)
+
+        x_tcn = self.channel_attn(x_tcn)
+        x_tcn = x_tcn.squeeze(1) if x_tcn.dim() == 3 and x_tcn.size(1) == 1 else x_tcn.mean(dim=-1)
+
+        x_trans = x_proj.unsqueeze(1)
+
+        attn_out1, _ = self.mha1(x_trans)
+        x_trans = self.layer_norm1(x_trans + attn_out1)
+
+        attn_out2, _ = self.mha2(x_trans)
+        x_trans = self.layer_norm2(x_trans + attn_out2)
+
+        ffn_out = self.ffn(x_trans)
+        x_trans = x_trans + ffn_out
+
+        lstm_out, _ = self.bilstm(x_trans)
+        lstm_out = lstm_out[:, -1, :]
+
+        combined = torch.cat([x_tcn, lstm_out], dim=1)
+
+        output = self.classifier(combined)
+
+        return output
+
+class FocalLoss(nn.Module):
+    def __init__(self, alpha=0.25, gamma=2):
+        super().__init__()
+        self.alpha = alpha
+        self.gamma = gamma
+
+    def forward(self, inputs, targets):
+        ce_loss = F.cross_entropy(inputs, targets, reduction='none')
+        pt = torch.exp(-ce_loss)
+        focal_loss = self.alpha * (1-pt)**self.gamma * ce_loss
+        return focal_loss.mean()
+
+def train_epoch(model, dataloader, criterion, optimizer, device):
+    model.train()
+    running_loss = 0.0
+    all_preds = []
+    all_labels = []
+
+    pbar = tqdm(dataloader, desc='Training')
+    for data, labels in pbar:
+        data, labels = data.to(device), labels.to(device)
+
+        optimizer.zero_grad()
+        outputs = model(data)
+        loss = criterion(outputs, labels)
+        loss.backward()
+
+        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
+
+        optimizer.step()
+
+        running_loss += loss.item()
+        _, preds = torch.max(outputs, 1)
+        all_preds.extend(preds.cpu().numpy())
+        all_labels.extend(labels.cpu().numpy())
+
+        pbar.set_postfix({'loss': loss.item()})
+
+    epoch_loss = running_loss / len(dataloader)
+    epoch_f1 = f1_score(all_labels, all_preds, average='weighted')
+
+    return epoch_loss, epoch_f1
+
+def validate(model, dataloader, criterion, device):
+    model.eval()
+    running_loss = 0.0
+    all_preds = []
+    all_labels = []
+    all_probs = []
+
+    with torch.no_grad():
+        for data, labels in tqdm(dataloader, desc='Validation'):
+            data, labels = data.to(device), labels.to(device)
+
+            outputs = model(data)
+            loss = criterion(outputs, labels)
+
+            running_loss += loss.item()
+            probs = F.softmax(outputs, dim=1)
+            _, preds = torch.max(outputs, 1)
+
+            all_preds.extend(preds.cpu().numpy())
+            all_labels.extend(labels.cpu().numpy())
+            all_probs.extend(probs.cpu().numpy()[:, 1])
+
+    epoch_loss = running_loss / len(dataloader)
+    epoch_f1 = f1_score(all_labels, all_preds, average='weighted')
+    epoch_auc = roc_auc_score(all_labels, all_probs)
+
+    return epoch_loss, epoch_f1, epoch_auc, all_preds, all_labels
+
+def main():
+    BATCH_SIZE = 64
+    LEARNING_RATE = 0.001
+    NUM_EPOCHS = 100
+    PATIENCE = 15
+
+    print("Loading datasets...")
+    train_dataset = EpilepsyDataset(r'train/path')
+    val_dataset = EpilepsyDataset(r'val/path')
+    test_dataset = EpilepsyDataset(r'test/path')
+
+    train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)
+    val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)
+    test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)
+
+    print(f"Train samples: {len(train_dataset)}")
+    print(f"Val samples: {len(val_dataset)}")
+    print(f"Test samples: {len(test_dataset)}")
+
+    model = AdvancedEpilepsyDetector(input_dim=178, num_classes=2, dropout=0.3).to(device)
+
+    total_params = sum(p.numel() for p in model.parameters())
+    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"\nTotal parameters: {total_params:,}")
+    print(f"Trainable parameters: {trainable_params:,}")
+
+    criterion = FocalLoss(alpha=0.25, gamma=2)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=0.01)
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5,
+                                                            patience=5, verbose=True)
+
+    best_val_f1 = 0.0
+    patience_counter = 0
+    train_losses, val_losses = [], []
+    train_f1s, val_f1s = [], []
+
+    print("\nStarting training...\n")
+
+    for epoch in range(NUM_EPOCHS):
+        print(f"Epoch {epoch+1}/{NUM_EPOCHS}")
+        print("-" * 50)
+
+        train_loss, train_f1 = train_epoch(model, train_loader, criterion, optimizer, device)
+
+        val_loss, val_f1, val_auc, _, _ = validate(model, val_loader, criterion, device)
+
+        scheduler.step(val_loss)
+
+        train_losses.append(train_loss)
+        val_losses.append(val_loss)
+        train_f1s.append(train_f1)
+        val_f1s.append(val_f1)
+
+        print(f"Train Loss: {train_loss:.4f} | Train F1: {train_f1:.4f}")
+        print(f"Val Loss: {val_loss:.4f} | Val F1: {val_f1:.4f} | Val AUC: {val_auc:.4f}\n")
+
+        if val_f1 > best_val_f1:
+            best_val_f1 = val_f1
+            patience_counter = 0
+            torch.save({
+                'epoch': epoch,
+                'model_state_dict': model.state_dict(),
+                'optimizer_state_dict': optimizer.state_dict(),
+                'val_f1': val_f1,
+                'val_auc': val_auc
+            }, r'best_epilepsy_model.pth')
+            print(f"[SAVED] Model saved with Val F1: {val_f1:.4f}\n")
+        else:
+            patience_counter += 1
+            if patience_counter >= PATIENCE:
+                print(f"\nEarly stopping triggered after {epoch+1} epochs")
+                break
+
+    print("\nLoading best model for testing...")
+    checkpoint = torch.load(r'best_epilepsy_model.pth')
+    model.load_state_dict(checkpoint['model_state_dict'])
+
+    print("\nEvaluating on test set...")
+    test_loss, test_f1, test_auc, test_preds, test_labels = validate(model, test_loader, criterion, device)
+
+    print(f"\nTest Results:")
+    print(f"Test Loss: {test_loss:.4f}")
+    print(f"Test F1: {test_f1:.4f}")
+    print(f"Test AUC: {test_auc:.4f}")
+
+    print("\nClassification Report:")
+    print(classification_report(test_labels, test_preds, target_names=['Non-Seizure', 'Seizure']))
+
+    # Confusion matrix
+    cm = confusion_matrix(test_labels, test_preds)
+    plt.figure(figsize=(10, 8))
+    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
+                xticklabels=['Non-Seizure', 'Seizure'],
+                yticklabels=['Non-Seizure', 'Seizure'])
+    plt.title('Confusion Matrix - Epilepsy Detection')
+    plt.ylabel('True Label')
+    plt.xlabel('Predicted Label')
+    plt.savefig(r'confusion_matrix.png', dpi=300, bbox_inches='tight')
+    plt.close()
+
+    plt.figure(figsize=(15, 5))
+
+    plt.subplot(1, 2, 1)
+    plt.plot(train_losses, label='Train Loss')
+    plt.plot(val_losses, label='Val Loss')
+    plt.xlabel('Epoch')
+    plt.ylabel('Loss')
+    plt.title('Training and Validation Loss')
+    plt.legend()
+    plt.grid(True)
+
+    plt.subplot(1, 2, 2)
+    plt.plot(train_f1s, label='Train F1')
+    plt.plot(val_f1s, label='Val F1')
+    plt.xlabel('Epoch')
+    plt.ylabel('F1 Score')
+    plt.title('Training and Validation F1 Score')
+    plt.legend()
+    plt.grid(True)
+
+    plt.savefig(r'training_curves.png', dpi=300, bbox_inches='tight')
+    plt.close()
+
+    print("\nTraining completed! Results saved.")
+
+if __name__ == "__main__":
+    main()