classifier.py

import os
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
from tqdm import tqdm
import json
import sys

# 添加路径
sys.path.append('/home/ldy/Closed_loop_optimizing')
sys.path.append('/home/ldy/Closed_loop_optimizing/model')

from model.ATMS_retrieval import ATMS, get_eeg_features
import open_clip

class EEGEmotionDataset(Dataset):
    """EEG情绪分类数据集"""
    
    def __init__(self, eeg_files, labels, transform=None):
        """
        Args:
            eeg_files: EEG文件路径列表
            labels: 对应的标签列表 (0或1)
            transform: 可选的数据变换
        """
        self.eeg_files = eeg_files
        self.labels = labels
        self.transform = transform
        
        assert len(eeg_files) == len(labels), "EEG文件数量与标签数量不匹配"
    
    def __len__(self):
        return len(self.eeg_files)
    
    def __getitem__(self, idx):
        # 加载EEG数据
        eeg_data = np.load(self.eeg_files[idx])
        
        # 确保数据格式正确 [channels, time_points]
        if eeg_data.ndim == 1:
            eeg_data = eeg_data.reshape(1, -1)
        
        if self.transform:
            eeg_data = self.transform(eeg_data)
        
        # 转换为tensor
        eeg_tensor = torch.FloatTensor(eeg_data)
        label = torch.LongTensor([self.labels[idx]])[0]
        
        return eeg_tensor, label

class EEGTransform:
    """EEG数据预处理变换"""
    
    def __init__(self, target_length=None, normalize=True):
        self.target_length = target_length
        self.normalize = normalize
    
    def __call__(self, eeg_data):
        # 时间维度裁剪或填充
        if self.target_length is not None:
            if eeg_data.shape[1] > self.target_length:
                eeg_data = eeg_data[:, :self.target_length]
            elif eeg_data.shape[1] < self.target_length:
                padding = np.zeros((eeg_data.shape[0], self.target_length - eeg_data.shape[1]))
                eeg_data = np.concatenate([eeg_data, padding], axis=1)
        
        # 标准化
        if self.normalize:
            eeg_data = (eeg_data - np.mean(eeg_data, axis=1, keepdims=True)) / (np.std(eeg_data, axis=1, keepdims=True) + 1e-8)
        
        return eeg_data

class EmotionClassifier(nn.Module):
    """基于ATM-S+CLIP的情绪分类器"""
    
    def __init__(self, atms_model, clip_model, feature_dim=1024, hidden_dim=512, num_classes=2, dropout=0.5):
        super(EmotionClassifier, self).__init__()
        
        self.atms_model = atms_model
        self.clip_model = clip_model
        
        # 冻结预训练模型参数
        for param in self.atms_model.parameters():
            param.requires_grad = False
        for param in self.clip_model.parameters():
            param.requires_grad = False
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.Linear(feature_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, num_classes)
        )
        
    def forward(self, eeg_data):
        """
        Args:
            eeg_data: EEG数据 [batch_size, channels, time_points]
        Returns:
            logits: 分类输出 [batch_size, num_classes]
        """
        # 使用ATM-S提取EEG特征
        with torch.no_grad():
            eeg_features = get_eeg_features(self.atms_model, eeg_data, eeg_data.device, 'sub-01')
        
        # 通过分类头
        logits = self.classifier(eeg_features)
        
        return logits

class EmotionClassifierTrainer:
    """情绪分类器训练器"""
    
    def __init__(self, model, device='cuda'):
        self.model = model.to(device)
        self.device = device
        self.train_losses = []
        self.val_losses = []
        self.train_accs = []
        self.val_accs = []
    
    def train_epoch(self, train_loader, optimizer, criterion):
        """训练一个epoch"""
        self.model.train()
        total_loss = 0
        correct = 0
        total = 0
        
        for batch_idx, (eeg_data, labels) in enumerate(tqdm(train_loader, desc="Training")):
            eeg_data, labels = eeg_data.to(self.device), labels.to(self.device)
            
            optimizer.zero_grad()
            outputs = self.model(eeg_data)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
            _, predicted = outputs.max(1)
            total += labels.size(0)
            correct += predicted.eq(labels).sum().item()
        
        avg_loss = total_loss / len(train_loader)
        accuracy = 100. * correct / total
        
        self.train_losses.append(avg_loss)
        self.train_accs.append(accuracy)
        
        return avg_loss, accuracy
    
    def validate_epoch(self, val_loader, criterion):
        """验证一个epoch"""
        self.model.eval()
        total_loss = 0
        correct = 0
        total = 0
        all_preds = []
        all_labels = []
        
        with torch.no_grad():
            for eeg_data, labels in tqdm(val_loader, desc="Validating"):
                eeg_data, labels = eeg_data.to(self.device), labels.to(self.device)
                
                outputs = self.model(eeg_data)
                loss = criterion(outputs, labels)
                
                total_loss += loss.item()
                _, predicted = outputs.max(1)
                total += labels.size(0)
                correct += predicted.eq(labels).sum().item()
                
                all_preds.extend(predicted.cpu().numpy())
                all_labels.extend(labels.cpu().numpy())
        
        avg_loss = total_loss / len(val_loader)
        accuracy = 100. * correct / total
        
        self.val_losses.append(avg_loss)
        self.val_accs.append(accuracy)
        
        return avg_loss, accuracy, all_preds, all_labels
    
    def train(self, train_loader, val_loader, num_epochs=50, lr=1e-3, weight_decay=1e-4, save_path=None):
        """完整训练流程"""
        criterion = nn.CrossEntropyLoss()
        optimizer = torch.optim.Adam(
            filter(lambda p: p.requires_grad, self.model.parameters()), 
            lr=lr, weight_decay=weight_decay
        )
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer, mode='max', factor=0.5, patience=5, verbose=True
        )
        
        best_val_acc = 0
        best_model_state = None
        
        print(f"开始训练，共{num_epochs}个epoch")
        print("-" * 60)
        
        for epoch in range(num_epochs):
            print(f"Epoch {epoch+1}/{num_epochs}")
            
            # 训练
            train_loss, train_acc = self.train_epoch(train_loader, optimizer, criterion)
            
            # 验证
            val_loss, val_acc, val_preds, val_labels = self.validate_epoch(val_loader, criterion)
            
            # 学习率调度
            scheduler.step(val_acc)
            
            print(f"Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.2f}%")
            print(f"Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.2f}%")
            print("-" * 60)
            
            # 保存最佳模型
            if val_acc > best_val_acc:
                best_val_acc = val_acc
                best_model_state = self.model.state_dict().copy()
                
                if save_path:
                    torch.save({
                        'model_state_dict': best_model_state,
                        'val_acc': best_val_acc,
                        'epoch': epoch,
                        'train_losses': self.train_losses,
                        'val_losses': self.val_losses,
                        'train_accs': self.train_accs,
                        'val_accs': self.val_accs
                    }, save_path)
                    print(f"保存最佳模型到: {save_path}")
        
        # 加载最佳模型
        if best_model_state:
            self.model.load_state_dict(best_model_state)
        
        return self.train_losses, self.val_losses, self.train_accs, self.val_accs
    
    def plot_training_history(self, save_path=None):
        """绘制训练历史"""
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
        
        # 损失曲线
        ax1.plot(self.train_losses, label='Training Loss', color='blue')
        ax1.plot(self.val_losses, label='Validation Loss', color='red')
        ax1.set_title('Training and Validation Loss')
        ax1.set_xlabel('Epoch')
        ax1.set_ylabel('Loss')
        ax1.legend()
        ax1.grid(True)
        
        # 准确率曲线
        ax2.plot(self.train_accs, label='Training Accuracy', color='blue')
        ax2.plot(self.val_accs, label='Validation Accuracy', color='red')
        ax2.set_title('Training and Validation Accuracy')
        ax2.set_xlabel('Epoch')
        ax2.set_ylabel('Accuracy (%)')
        ax2.legend()
        ax2.grid(True)
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"训练历史图保存到: {save_path}")
        
        plt.show()

def evaluate_model(model, test_loader, device, class_names=['Negative', 'Positive']):
    """评估模型性能"""
    model.eval()
    all_preds = []
    all_labels = []
    all_probs = []
    
    with torch.no_grad():
        for eeg_data, labels in tqdm(test_loader, desc="Testing"):
            eeg_data, labels = eeg_data.to(device), labels.to(device)
            
            outputs = model(eeg_data)
            probs = F.softmax(outputs, dim=1)
            _, predicted = outputs.max(1)
            
            all_preds.extend(predicted.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())
            all_probs.extend(probs.cpu().numpy())
    
    # 计算指标
    accuracy = accuracy_score(all_labels, all_preds)
    report = classification_report(all_labels, all_preds, target_names=class_names)
    cm = confusion_matrix(all_labels, all_preds)
    
    print(f"测试准确率: {accuracy:.4f}")
    print("\n分类报告:")
    print(report)
    
    # 绘制混淆矩阵
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=class_names, yticklabels=class_names)
    plt.title('Confusion Matrix')
    plt.ylabel('True Label')
    plt.xlabel('Predicted Label')
    plt.show()
    
    return accuracy, all_preds, all_labels, all_probs

def load_emotion_data(data_dir, label_file):
    """
    加载情绪数据
    
    Args:
        data_dir: EEG文件目录
        label_file: 标签文件路径 (JSON格式: {"file1.npy": 0, "file2.npy": 1, ...})
    
    Returns:
        eeg_files: EEG文件路径列表
        labels: 对应标签列表
    """
    # 加载标签
    with open(label_file, 'r') as f:
        label_dict = json.load(f)
    
    eeg_files = []
    labels = []
    
    for filename, label in label_dict.items():
        file_path = os.path.join(data_dir, filename)
        if os.path.exists(file_path):
            eeg_files.append(file_path)
            labels.append(label)
        else:
            print(f"警告: 文件不存在 {file_path}")
    
    print(f"加载了 {len(eeg_files)} 个EEG文件")
    print(f"标签分布: {np.bincount(labels)}")
    
    return eeg_files, labels

def main():
    """主函数"""
    # 设备配置
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"使用设备: {device}")
    
    # 路径配置
    sub = 'sub-01'
    data_dir = "/path/to/your/eeg/data"  # 替换为您的EEG数据目录
    label_file = "/path/to/your/labels.json"  # 替换为您的标签文件
    save_dir = "/home/ldy/Closed_loop_optimizing/ml-project/emotion_classifier"
    os.makedirs(save_dir, exist_ok=True)
    
    # 加载预训练模型
    print("加载预训练模型...")
    
    # 加载ATM-S模型
    f_encoder = f"/mnt/dataset0/kyw/closed-loop/sub_model/{sub}/diffusion_alexnet/pretrained_True/gene_gene/ATM_S_reconstruction_scale_0_1000_40.pth"
    checkpoint = torch.load(f_encoder, map_location=device)
    atms_model = ATMS()
    atms_model.load_state_dict(checkpoint['eeg_model_state_dict'])
    atms_model.to(device)
    
    # 加载CLIP模型
    model_type = 'ViT-H-14'
    clip_model, _, _ = open_clip.create_model_and_transforms(
        model_type, pretrained='laion2b_s32b_b79k', precision='fp32', device=device
    )
    
    # 加载数据
    print("加载数据...")
    eeg_files, labels = load_emotion_data(data_dir, label_file)
    
    # 数据划分
    train_files, temp_files, train_labels, temp_labels = train_test_split(
        eeg_files, labels, test_size=0.4, random_state=42, stratify=labels
    )
    val_files, test_files, val_labels, test_labels = train_test_split(
        temp_files, temp_labels, test_size=0.5, random_state=42, stratify=temp_labels
    )
    
    print(f"训练集: {len(train_files)} 样本")
    print(f"验证集: {len(val_files)} 样本")
    print(f"测试集: {len(test_files)} 样本")
    
    # 创建数据集和数据加载器
    transform = EEGTransform(target_length=250, normalize=True)  # 1秒数据，250Hz采样率
    
    train_dataset = EEGEmotionDataset(train_files, train_labels, transform=transform)
    val_dataset = EEGEmotionDataset(val_files, val_labels, transform=transform)
    test_dataset = EEGEmotionDataset(test_files, test_labels, transform=transform)
    
    train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=4)
    val_loader = DataLoader(val_dataset, batch_size=16, shuffle=False, num_workers=4)
    test_loader = DataLoader(test_dataset, batch_size=16, shuffle=False, num_workers=4)
    
    # 创建分类器
    print("创建分类器...")
    classifier = EmotionClassifier(
        atms_model=atms_model,
        clip_model=clip_model,
        feature_dim=1024,
        hidden_dim=512,
        num_classes=2,
        dropout=0.5
    )
    
    # 训练
    print("开始训练...")
    trainer = EmotionClassifierTrainer(classifier, device=device)
    model_save_path = os.path.join(save_dir, 'best_emotion_classifier.pth')
    
    train_losses, val_losses, train_accs, val_accs = trainer.train(
        train_loader=train_loader,
        val_loader=val_loader,
        num_epochs=50,
        lr=1e-3,
        weight_decay=1e-4,
        save_path=model_save_path
    )
    
    # 绘制训练历史
    history_plot_path = os.path.join(save_dir, 'training_history.png')
    trainer.plot_training_history(save_path=history_plot_path)
    
    # 测试评估
    print("评估模型...")
    accuracy, preds, labels, probs = evaluate_model(
        classifier, test_loader, device, class_names=['Negative', 'Positive']
    )
    
    # 保存结果
    results = {
        'test_accuracy': accuracy,
        'train_losses': train_losses,
        'val_losses': val_losses,
        'train_accs': train_accs,
        'val_accs': val_accs
    }
    
    with open(os.path.join(save_dir, 'results.json'), 'w') as f:
        json.dump(results, f, indent=2)
    
    print(f"训练完成！最终测试准确率: {accuracy:.4f}")
    print(f"模型和结果保存到: {save_dir}")

if __name__ == "__main__":
    main()