evaluate.py

# -*- coding: utf-8 -*-
"""
评估脚本 - 模型评估与性能分析
"""
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from sklearn.metrics import confusion_matrix, classification_report, f1_score
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
import seaborn as sns

from config import config
from dataset import generate_dataset
from model_v2 import DualBranchClassifier


def normalize_data(X):
    """数据归一化，同时计算FFT作为第三通道"""
    # 按样本归一化（对I和Q）
    mean = X.mean(axis=(1, 2), keepdims=True)
    std = X.std(axis=(1, 2), keepdims=True)
    std[std == 0] = 1
    X_norm = (X - mean) / std
    
    # 计算FFT幅度作为第三通道
    batch_size = X.shape[0]
    signal_length = X.shape[2]
    fft_features = np.zeros((batch_size, 1, signal_length // 2), dtype=np.float32)
    
    for i in range(batch_size):
        complex_signal = X[i, 0, :] + 1j * X[i, 1, :]
        fft_result = np.fft.fft(complex_signal)
        fft_magnitude = np.abs(fft_result[:signal_length // 2])
        fft_magnitude = fft_magnitude / (np.max(fft_magnitude) + 1e-8)
        fft_features[i, 0, :] = fft_magnitude
    
    # 拼接：I、Q、FFT幅度 -> (batch, 3, signal_length//2)
    X_3ch = np.concatenate([X_norm[:, :, :signal_length // 2], fft_features], axis=1)
    
    return X_3ch


def load_model(model_path='./checkpoints/best_model.pth'):
    """加载训练好的模型"""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = DualBranchClassifier(num_classes=config.NUM_CLASSES)
    
    checkpoint = torch.load(model_path, map_location=device)
    model.load_state_dict(checkpoint['model_state_dict'])
    model = model.to(device)
    model.eval()
    
    print(f"模型加载成功！验证集准确率: {checkpoint['val_acc']:.2f}%")
    return model, device


def evaluate_model(model, X, y, device):
    """完整模型评估"""
    X = normalize_data(X)
    
    # 划分数据集
    X_train, X_temp, y_train, y_temp = train_test_split(
        X, y, test_size=config.VAL_RATIO + config.TEST_RATIO,
        random_state=config.RANDOM_SEED, stratify=y
    )
    val_size = config.TEST_RATIO / (config.VAL_RATIO + config.TEST_RATIO)
    X_val, X_test, y_val, y_test = train_test_split(
        X_temp, y_temp, test_size=val_size,
        random_state=config.RANDOM_SEED, stratify=y_temp
    )
    
    # 测试集评估
    test_dataset = TensorDataset(torch.FloatTensor(X_test), torch.LongTensor(y_test))
    test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)
    
    all_preds = []
    all_labels = []
    
    with torch.no_grad():
        for inputs, labels in test_loader:
            inputs = inputs.to(device)
            outputs = model(inputs)
            _, predicted = outputs.max(1)
            all_preds.extend(predicted.cpu().numpy())
            all_labels.extend(labels.numpy())
    
    all_preds = np.array(all_preds)
    all_labels = np.array(all_labels)
    
    # 计算指标
    accuracy = 100 * np.mean(all_preds == all_labels)
    f1 = f1_score(all_labels, all_preds, average='weighted')
    
    print(f"\n{'='*50}")
    print(f"测试集评估结果")
    print(f"{'='*50}")
    print(f"准确率: {accuracy:.2f}%")
    print(f"Weighted F1 Score: {f1:.4f}")
    print(f"\n分类报告:")
    print(classification_report(all_labels, all_preds, 
                                target_names=config.MODULATION_TYPES))
    
    return all_preds, all_labels, accuracy, f1


def evaluate_by_snr(model, X, y, snrs, device):
    """按信噪比分组的评估"""
    X = normalize_data(X)
    
    snr_acc = {}
    
    for snr in config.SNR_RANGE:
        mask = snrs == snr
        X_snr = X[mask]
        y_snr = y[mask]
        
        if len(X_snr) == 0:
            continue
        
        test_dataset = TensorDataset(torch.FloatTensor(X_snr), torch.LongTensor(y_snr))
        test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)
        
        correct = 0
        total = 0
        
        with torch.no_grad():
            for inputs, labels in test_loader:
                inputs = inputs.to(device)
                outputs = model(inputs)
                _, predicted = outputs.max(1)
                total += labels.size(0)
                correct += predicted.eq(labels).sum().item()
        
        acc = 100.0 * correct / total
        snr_acc[snr] = acc
        print(f"SNR = {snr:3d} dB: 准确率 = {acc:.2f}%")
    
    return snr_acc


def plot_confusion_matrix(y_true, y_pred, save_path='./results/confusion_matrix.png'):
    """绘制混淆矩阵"""
    cm = confusion_matrix(y_true, y_pred)
    
    plt.figure(figsize=(12, 10))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                xticklabels=config.MODULATION_TYPES,
                yticklabels=config.MODULATION_TYPES)
    plt.title('混淆矩阵 - 调制类型识别', fontsize=14)
    plt.xlabel('预测类别', fontsize=12)
    plt.ylabel('真实类别', fontsize=12)
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"\n混淆矩阵已保存到: {save_path}")


def plot_snr_accuracy(snr_acc, save_path='./results/snr_accuracy.png'):
    """绘制不同信噪比下的准确率曲线"""
    snrs = sorted(snr_acc.keys())
    accs = [snr_acc[s] for s in snrs]
    
    plt.figure(figsize=(10, 6))
    plt.plot(snrs, accs, 'bo-', linewidth=2, markersize=8)
    plt.fill_between(snrs, accs, alpha=0.3)
    
    plt.xlabel('信噪比 (dB)', fontsize=12)
    plt.ylabel('准确率 (%)', fontsize=12)
    plt.title('不同信噪比下的分类准确率', fontsize=14)
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.xticks(snrs)
    plt.ylim(0, 105)
    
    for s, a in zip(snrs, accs):
        plt.annotate(f'{a:.1f}%', (s, a), textcoords="offset points", 
                    xytext=(0, 10), ha='center', fontsize=9)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"SNR-准确率曲线已保存到: {save_path}")


def plot_accuracy_by_class(y_true, y_pred, save_path='./results/class_accuracy.png'):
    """绘制各类别的准确率"""
    cm = confusion_matrix(y_true, y_pred)
    class_acc = cm.diagonal() / cm.sum(axis=1)
    
    plt.figure(figsize=(12, 5))
    colors = plt.cm.RdYlGn(class_acc)
    bars = plt.bar(config.MODULATION_TYPES, class_acc * 100, color=colors)
    
    plt.xlabel('调制类型', fontsize=12)
    plt.ylabel('准确率 (%)', fontsize=12)
    plt.title('各类别的分类准确率', fontsize=14)
    plt.xticks(rotation=45, ha='right')
    plt.ylim(0, 105)
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    
    for bar, acc in zip(bars, class_acc):
        plt.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 1,
                f'{acc*100:.1f}%', ha='center', va='bottom', fontsize=9)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches='tight')
    plt.close()
    print(f"类别准确率图已保存到: {save_path}")


def main():
    """主函数"""
    # 加载模型
    model, device = load_model()
    
    # 生成数据
    print("正在生成测试数据...")
    X, y, snrs = generate_dataset()
    
    # 整体评估
    y_pred, y_true, accuracy, f1 = evaluate_model(model, X, y, device)
    
    # 按SNR评估
    print(f"\n{'='*50}")
    print("按信噪比分组的评估结果")
    print(f"{'='*50}")
    snr_acc = evaluate_by_snr(model, X, y, snrs, device)
    
    # 生成可视化
    plot_confusion_matrix(y_true, y_pred)
    plot_snr_accuracy(snr_acc)
    plot_accuracy_by_class(y_true, y_pred)
    
    print(f"\n所有结果已保存到 ./results/ 目录")


if __name__ == '__main__':
    main()