Upload 4 files

EpilepsyNet.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f39a4c7a0ce846a6d57977917288297090c0a73a77acd5fcbfab2ab460bbf6de
+size 3617604

EpilepsyNet_model.py

@@ -0,0 +1,278 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+
+
+
+def extract_upper_triangle(corr_matrices):
+    """
+    Extract upper triangles from correlation matrices
+    
+    Args:
+        corr_matrices: numpy array of shape (n_segments, n_channels, n_channels)
+        
+    Returns:
+        numpy array of shape (n_segments, n_features) where n_features = n_channels*(n_channels-1)/2
+    """
+    n_segments, n_channels, _ = corr_matrices.shape
+    n_features = n_channels * (n_channels - 1) // 2
+    
+    flattened = np.zeros((n_segments, n_features))
+    
+    for i in range(n_segments):
+        # Get upper triangle indices (excluding diagonal)
+        upper_indices = np.triu_indices(n_channels, k=1)
+        # Extract values
+        flattened[i] = corr_matrices[i][upper_indices]
+    
+    return flattened
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, embed_dim, num_heads, dropout=0.3):
+        super(MultiHeadAttention, self).__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        assert self.head_dim * num_heads == embed_dim, "embed_dim must be divisible by num_heads"
+        
+        # Linear projections for Q, K, V
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        
+        # Final projection after concatenating heads
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+        
+        # Dropout
+        self.dropout = nn.Dropout(dropout)
+        
+        # Softmax for attention weights
+        self.softmax = nn.Softmax(dim=-1)
+        
+    def forward(self, x, mask=None):
+        batch_size = x.size(0)
+        
+        # Project Q, K, V
+        Q = self.q_proj(x)  # (batch_size, seq_len, embed_dim)
+        K = self.k_proj(x)    # (batch_size, seq_len, embed_dim)
+        V = self.v_proj(x)  # (batch_size, seq_len, embed_dim)
+        
+        # Split into multiple heads
+        Q = Q.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)  # (batch_size, num_heads, seq_len, head_dim)
+        K = K.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)  # (batch_size, num_heads, seq_len, head_dim)
+        V = V.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)  # (batch_size, num_heads, seq_len, head_dim)
+        
+        # Calculate attention scores
+        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)  # (batch_size, num_heads, seq_len, seq_len)
+        
+        # Apply mask (if provided)
+        if mask is not None:
+            scores = scores.masked_fill(mask == 0, float('-inf'))
+            
+        # Apply softmax to get attention weights
+        attn_weights = self.softmax(scores)  # (batch_size, num_heads, seq_len, seq_len)
+        attn_weights = self.dropout(attn_weights)
+        
+        # Calculate weighted output
+        attn_output = torch.matmul(attn_weights, V)  # (batch_size, num_heads, seq_len, head_dim)
+        
+        # Recompose heads
+        attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, -1, self.embed_dim)  # (batch_size, seq_len, embed_dim)
+        
+        # Pass through final projection
+        output = self.out_proj(attn_output)  # (batch_size, seq_len, embed_dim)
+        
+        return output, attn_weights
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, embed_dim, max_seq_length=100):
+        super(PositionalEncoding, self).__init__()
+        
+        # Create positional encoding matrix
+        pe = torch.zeros(max_seq_length, embed_dim)
+        position = torch.arange(0, max_seq_length, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, embed_dim, 2).float() * (-np.log(10000.0) / embed_dim))
+        
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        
+        # Register as buffer (not a parameter)
+        self.register_buffer('pe', pe.unsqueeze(0))
+        
+    def forward(self, x):
+        # Add positional encoding to input
+        # x: [batch_size, seq_len, embed_dim]
+        return x + self.pe[:, :x.size(1)]
+
+class TimeSeriesAttentionClassifier(nn.Module):
+    def __init__(self, input_dim, embed_dim, num_heads, num_classes=2, dropout=0.2):
+        super(TimeSeriesAttentionClassifier, self).__init__()
+        
+        # Project flattened correlation features to embedding space
+        self.embedding = nn.Linear(input_dim, embed_dim)
+        
+        # Positional encoding
+        self.pos_encoding = PositionalEncoding(embed_dim)
+        
+        # Multi-head attention
+        self.attention = MultiHeadAttention(embed_dim, num_heads, dropout)
+        
+        # Layer normalization
+        self.layer_norm1 = nn.LayerNorm(embed_dim)
+        self.layer_norm2 = nn.LayerNorm(embed_dim)
+        
+        # Feed-forward network
+        self.ffn = nn.Sequential(
+            nn.Linear(embed_dim, embed_dim * 4),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(embed_dim * 4, embed_dim)
+        )
+        
+        # Output layer
+        self.classifier = nn.Sequential(
+            nn.Linear(embed_dim, embed_dim // 2),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(embed_dim // 2, 1),
+            nn.Sigmoid()
+        )
+        
+    def forward(self, x):
+        # batch_size, seq_len, input_dim = x.shape
+        
+        # Project to embedding space
+        x = self.embedding(x)
+        
+        # Add positional encoding
+        x = self.pos_encoding(x)
+        
+        # Self-attention (use x for query, key, and value)
+        residual = x
+        x, attention_weights = self.attention(x)
+        x = self.layer_norm1(x + residual)
+        
+        # Feed-forward network with residual connection
+        residual = x
+        x = self.ffn(x)
+        x = self.layer_norm2(x + residual)
+        
+        # Global average pooling over sequence dimension
+        x = torch.mean(x, dim=1)
+        
+        # Classification
+        logits = self.classifier(x)
+        
+        return logits, attention_weights
+
+def train_model(model, train_loader, val_loader, num_epochs=50, learning_rate=1e-4, weight_decay=1e-5,  patience=10, scheduler_factor=0.5, min_lr=1e-6):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = model.to(device)
+    
+    # Changed from CrossEntropyLoss to BCELoss for binary classification with sigmoid
+    criterion = nn.BCELoss()
+    # optimizer = optim.Adam(model.parameters(), lr=learning_rate)
+    
+    # Add L2 regularization through weight_decay parameter in Adam
+    optimizer = optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
+    
+    # Learning rate scheduler - reduce LR when validation loss plateaus
+    scheduler = optim.lr_scheduler.ReduceLROnPlateau(
+        optimizer, 
+        mode='min',
+        factor=scheduler_factor,
+        patience=patience,
+        verbose=True,
+        min_lr=min_lr
+    )
+
+    train_losses = []
+    val_losses = []
+    val_accuracies = []
+    
+    
+    # Track best model and early stopping
+    best_val_loss = float('inf')
+    best_model_state = None
+    early_stop_counter = 0
+    early_stop_patience = patience * 2  # Stop after 2x the scheduler patience
+    
+    for epoch in range(num_epochs):
+        # Training
+        model.train()
+        train_loss = 0.0
+        for inputs, labels in train_loader:
+            inputs, labels = inputs.to(device), labels.to(device)
+            
+            # Convert labels to float and reshape for BCE loss
+            labels = labels.float().view(-1, 1)
+            
+            optimizer.zero_grad()
+            outputs, _ = model(inputs)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+            train_loss += loss.item()
+        
+        train_loss /= len(train_loader)
+        train_losses.append(train_loss)
+        
+        # Validation
+        model.eval()
+        val_loss = 0.0
+        correct = 0
+        total = 0
+        
+        with torch.no_grad():
+            for inputs, labels in val_loader:
+                inputs, labels = inputs.to(device), labels.to(device)
+                # Convert labels to float and reshape for BCE loss
+                labels = labels.float().view(-1, 1)
+                
+                outputs, _ = model(inputs)
+                loss = criterion(outputs, labels)
+                val_loss += loss.item()
+                
+                # For binary classification with sigmoid, prediction is 1 if output > 0.5
+                predicted = (outputs > 0.5).float()
+                total += labels.size(0)
+                correct += (predicted == labels).sum().item()
+        
+        val_loss /= len(val_loader)
+        val_losses.append(val_loss)
+        
+        accuracy = 100 * correct / total
+        val_accuracies.append(accuracy)
+        
+        # Learning rate scheduler step based on validation loss
+        scheduler.step(val_loss)
+        
+        # Print current learning rate
+        current_lr = optimizer.param_groups[0]['lr']
+
+
+        # Print epoch results
+        print(f'Epoch {epoch+1}/{num_epochs}, LR: {current_lr:.6f}, Train Loss: {train_loss:.4f}, '
+              f'Val Loss: {val_loss:.4f}, Val Accuracy: {accuracy:.2f}%')
+        
+        # Save best model
+        if val_loss < best_val_loss:
+            best_val_loss = val_loss
+            best_model_state = model.state_dict().copy()
+            early_stop_counter = 0
+        else:
+            early_stop_counter += 1
+        
+        # Early stopping
+        if early_stop_counter >= early_stop_patience:
+            print(f"Early stopping triggered after {epoch+1} epochs")
+            break
+
+    # Load best model weights
+    if best_model_state is not None:
+        model.load_state_dict(best_model_state)
+        print(f"Loaded best model with validation loss: {best_val_loss:.4f}")
+    
+    return train_losses, val_losses, val_accuracies
+

app.py

@@ -0,0 +1,69 @@
+import os
+import tempfile
+import uvicorn
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi.responses import JSONResponse
+
+# Import your prediction functions
+from prediction import predict_eeg_recording, predict_ensemble_eeg_recording
+
+app = FastAPI(title="EEG Epilepsy Prediction API")
+
+@app.get("/", tags=["Introduction Endpoints"])
+async def index():
+    """
+    Simply returns a welcome message!
+    """
+    message = (
+        "Hello world! Welcome to the EEG Epilepsy Prediction API. "
+        "Submit an EEG recording EDF file to the `/predict` endpoint to receive a prediction."
+    )
+    return message
+
+@app.post("/predict", tags=["Machine Learning"])
+async def predict_endpoint(
+    file: UploadFile = File(...),
+    model_choice: str = "2DCNN",
+    ensemble_method: str = None
+):
+    """
+    
+    Query parameters:
+      - model_choice: Choose one model among "2DCNN", "EEGNet", "EpilepsyNet", or "ensemble".
+      - ensemble_method: (Optional, required if model_choice is "ensemble") 
+                           The ensemble method to use ("average" or "voting").
+
+    """
+    print("Saving uploaded file as temporary file...")
+    try:
+        suffix = os.path.splitext(file.filename)[1]
+        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
+            tmp.write(await file.read())
+            tmp_path = tmp.name
+    except Exception as e:
+        raise HTTPException(status_code=500, detail="Error saving temporary file")
+    
+    print("Performing prediction using model_choice =", model_choice)
+    try:
+        if model_choice.lower() == "ensemble":
+            if ensemble_method is None:
+                raise HTTPException(status_code=400, detail="ensemble_method must be specified when using ensemble model_choice")
+            pred_label, mean_prob = predict_ensemble_eeg_recording(tmp_path, ensemble_method=ensemble_method, threshold=0.5)
+        else:
+            pred_label, mean_prob = predict_eeg_recording(tmp_path, model_name=model_choice, threshold=0.5)
+    except Exception as e:
+        os.remove(tmp_path)
+        raise HTTPException(status_code=400, detail=f"Prediction failed: {e}")
+    
+    os.remove(tmp_path)
+    
+    
+    response = {
+        "prediction": "epilepsy" if pred_label == 1 else "no epilepsy",
+        "confidence": mean_prob
+    }
+    print("Prediction complete, returning response...")
+    return JSONResponse(content=response)
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)

eegnet_model.py

@@ -0,0 +1,50 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+
+class EEGNet(nn.Module):
+    def __init__(self, n_channels=21, n_samples=1250, num_classes=2, dropout_rate=0.5):
+        super(EEGNet, self).__init__()
+
+        # Temporal convolution: learn temporal filters across time dimension
+        self.firstconv = nn.Sequential(
+            nn.Conv2d(1, 8, kernel_size=(1, 64), padding=(0, 32), bias=False),  # shape: (B, 8, C, T)
+            nn.BatchNorm2d(8)
+        )
+
+        # Depthwise spatial convolution: one spatial filter per temporal filter
+        self.depthwiseConv = nn.Sequential(
+            nn.Conv2d(8, 16, kernel_size=(n_channels, 1), groups=8, bias=False),  # shape: (B, 16, 1, T)
+            nn.BatchNorm2d(16),
+            nn.ELU(),
+            nn.AvgPool2d(kernel_size=(1, 4)),
+            nn.Dropout(dropout_rate)
+        )
+
+        # Separable convolution: combines temporal filters again
+        self.separableConv = nn.Sequential(
+            nn.Conv2d(16, 16, kernel_size=(1, 16), padding=(0, 8), bias=False),
+            nn.BatchNorm2d(16),
+            nn.ELU(),
+            nn.AvgPool2d(kernel_size=(1, 8)),
+            nn.Dropout(dropout_rate)
+        )
+
+        # Dynamically compute the flattened feature size after conv layers
+        dummy_input = torch.zeros(1, 1, n_channels, n_samples)
+        with torch.no_grad():
+            x = self.firstconv(dummy_input)
+            x = self.depthwiseConv(x)
+            x = self.separableConv(x)
+            flattened_size = x.reshape(1, -1).shape[1]  # dynamically computed
+
+        # Final classification layer
+        self.classifier = nn.Linear(flattened_size, num_classes)
+
+    def forward(self, x):
+        x = self.firstconv(x)
+        x = self.depthwiseConv(x)
+        x = self.separableConv(x)
+        x = x.reshape(x.size(0), -1)  # flatten
+        x = self.classifier(x)
+        return x