plot_sci_figures.py

#!/usr/bin/env python3
"""
SCI 论文级图表绘制脚本
=====================
用法:
  1. 训练后自动调用 (train.py 中已集成)
  2. 独立运行: python plot_sci_figures.py --result_dir outputs/kfold_train_xxx

生成图表:
  fig1_kfold_accuracy.png    — K折交叉验证准确率柱状图 (mean±std)
  fig2_learning_curves.png   — 各折训练/验证学习曲线 (如有epoch_history)
  fig3_confusion_matrix.png  — 混淆矩阵热力图 (如有evaluation_metrics)
  fig4_roc_curves.png        — 多类别ROC曲线 (如有roc_data)
  fig5_tsne_embedding.png    — t-SNE嵌入可视化 (如有embedding_data)
"""

import os
import sys
import json
import argparse
import numpy as np

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib import rcParams

# ============================================================
# SCI 绘图全局样式
# ============================================================

def set_sci_style():
    """设置 SCI 论文级绘图样式"""
    rcParams.update({
        # 字体: Times New Roman, 兼容中文
        'font.family': 'serif',
        'font.serif': ['Times New Roman', 'SimSun', 'DejaVu Serif'],
        'font.sans-serif': ['Arial', 'SimHei', 'DejaVu Sans'],
        'font.size': 10,
        'axes.labelsize': 11,
        'axes.titlesize': 12,
        'xtick.labelsize': 9,
        'ytick.labelsize': 9,
        'legend.fontsize': 9,

        # 坐标轴
        'axes.linewidth': 0.8,
        'axes.edgecolor': 'black',
        'axes.labelcolor': 'black',
        'axes.unicode_minus': False,
        'axes.spines.top': False,
        'axes.spines.right': False,

        # 刻度
        'xtick.direction': 'in',
        'ytick.direction': 'in',
        'xtick.major.size': 4,
        'ytick.major.size': 4,
        'xtick.major.width': 0.8,
        'ytick.major.width': 0.8,
        'xtick.minor.visible': True,
        'ytick.minor.visible': True,
        'xtick.minor.size': 2,
        'ytick.minor.size': 2,

        # 图例
        'legend.frameon': True,
        'legend.framealpha': 0.9,
        'legend.edgecolor': '0.8',
        'legend.borderpad': 0.4,

        # 网格
        'axes.grid': False,

        # 保存
        'savefig.dpi': 600,
        'savefig.bbox': 'tight',
        'savefig.pad_inches': 0.05,
        'figure.facecolor': 'white',
    })


# 配色方案 (色盲友好, 适配灰度打印)
COLORS = {
    'primary':   '#2166AC',   # 深蓝
    'secondary': '#B2182B',   # 深红
    'tertiary':  '#4DAF4A',   # 绿
    'quaternary':'#FF7F00',   # 橙
    'light_blue':'#92C5DE',
    'light_red': '#F4A582',
    'gray':      '#666666',
    'light_gray':'#CCCCCC',
}

FOLD_COLORS = ['#2166AC', '#B2182B', '#4DAF4A', '#FF7F00', '#984EA3']


# ============================================================
# Fig 1: K折交叉验证准确率柱状图
# ============================================================

def plot_kfold_accuracy(result_dir, summary, save_path):
    """
    绘制K折交叉验证准确率柱状图
    每个柱子 = 该折最佳验证准确率, 叠加 mean±std 标注
    """
    fold_results = summary['fold_results']
    fold_labels = [f"Fold {r['fold']}" for r in fold_results]
    accuracies = [r['best_val_acc'] for r in fold_results]

    overall = summary.get('overall_stats', {})
    mean_acc = overall.get('mean_accuracy', np.mean(accuracies))
    std_acc = overall.get('std_accuracy', np.std(accuracies))

    fig, ax = plt.subplots(figsize=(3.5, 3.0))

    x = np.arange(len(fold_labels))
    bars = ax.bar(x, accuracies, width=0.55, color=FOLD_COLORS[:len(fold_labels)],
                  edgecolor='black', linewidth=0.6, alpha=0.85, zorder=3)

    # 均值线
    ax.axhline(y=mean_acc, color=COLORS['secondary'], linestyle='--', linewidth=1.0,
               label=f'Mean = {mean_acc:.2f}%', zorder=4)

    # ±std 阴影带
    ax.axhspan(mean_acc - std_acc, mean_acc + std_acc,
               alpha=0.12, color=COLORS['secondary'], zorder=1)

    # 柱子顶部标注数值
    for bar, acc in zip(bars, accuracies):
        ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.3,
                f'{acc:.1f}%', ha='center', va='bottom', fontsize=8, fontweight='bold')

    # 标注 best epoch
    for i, r in enumerate(fold_results):
        best_ep = r.get('best_epoch', '?')
        ax.text(bars[i].get_x() + bars[i].get_width()/2, 5,
                f'Ep.{best_ep}', ha='center', va='bottom', fontsize=7,
                color='white', fontweight='bold')

    ax.set_xlabel('Cross-Validation Fold')
    ax.set_ylabel('Best Validation Accuracy (%)')
    ax.set_xticks(x)
    ax.set_xticklabels(fold_labels)
    ax.set_ylim(85, 102)
    ax.legend(loc='lower right')

    # 右上角标注 mean±std
    ax.text(0.98, 0.98, f'{mean_acc:.2f} $\\pm$ {std_acc:.2f}%',
            transform=ax.transAxes, ha='right', va='top',
            fontsize=10, fontweight='bold',
            bbox=dict(boxstyle='round,pad=0.3', facecolor='lightyellow',
                      edgecolor='gray', alpha=0.9))

    fig.tight_layout()
    fig.savefig(save_path)
    plt.close(fig)
    print(f"  [Fig1] K-fold accuracy -> {save_path}")


# ============================================================
# Fig 2: 学习曲线 (各折 train/val loss & accuracy)
# ============================================================

def plot_learning_curves(result_dir, summary, save_path):
    """
    绘制各折的训练/验证 loss 和 accuracy 学习曲线
    需要 epoch_history 数据
    """
    fold_results = summary['fold_results']

    # 检查是否有 epoch_history
    has_history = any('epoch_history' in r and len(r.get('epoch_history', [])) > 0
                      for r in fold_results)
    if not has_history:
        print("  [Fig2] 跳过: 无 epoch_history 数据 (需重新训练)")
        return

    fig, axes = plt.subplots(1, 2, figsize=(7.0, 3.0))

    # 左图: Loss
    ax_loss = axes[0]
    # 右图: Accuracy
    ax_acc = axes[1]

    for i, r in enumerate(fold_results):
        history = r.get('epoch_history', [])
        if not history:
            continue

        epochs = [h['epoch'] for h in history]
        train_loss = [h['train_loss'] for h in history]
        val_loss = [h['val_loss'] for h in history]
        train_acc = [h['train_acc'] for h in history]
        val_acc = [h['val_acc'] for h in history]

        color = FOLD_COLORS[i % len(FOLD_COLORS)]
        fold_label = f"Fold {r['fold']}"

        # Loss 曲线
        ax_loss.plot(epochs, train_loss, color=color, linestyle='-', linewidth=0.8,
                     alpha=0.5, label=f'{fold_label} Train' if i < 3 else None)
        ax_loss.plot(epochs, val_loss, color=color, linestyle='-', linewidth=1.3,
                     label=f'{fold_label} Val')

        # Accuracy 曲线
        ax_acc.plot(epochs, train_acc, color=color, linestyle='-', linewidth=0.8,
                    alpha=0.5, label=f'{fold_label} Train' if i < 3 else None)
        ax_acc.plot(epochs, val_acc, color=color, linestyle='-', linewidth=1.3,
                    label=f'{fold_label} Val')

        # 标记最佳 epoch
        best_ep = r.get('best_epoch', 0)
        if best_ep > 0 and best_ep <= len(val_acc):
            ax_acc.plot(best_ep, val_acc[best_ep-1], marker='*', color=color,
                        markersize=10, zorder=5)

    ax_loss.set_xlabel('Epoch')
    ax_loss.set_ylabel('Loss')
    ax_loss.set_title('(a) Training & Validation Loss')
    ax_loss.legend(loc='upper right', ncol=1, fontsize=7)

    ax_acc.set_xlabel('Epoch')
    ax_acc.set_ylabel('Accuracy (%)')
    ax_acc.set_title('(b) Training & Validation Accuracy')
    ax_acc.legend(loc='lower right', ncol=1, fontsize=7)
    ax_acc.set_ylim(80, 102)

    fig.tight_layout(w_pad=3.0)
    fig.savefig(save_path)
    plt.close(fig)
    print(f"  [Fig2] Learning curves -> {save_path}")


# ============================================================
# Fig 3: 混淆矩阵
# ============================================================

def plot_confusion_matrix(result_dir, summary, save_path):
    """
    绘制归一化混淆矩阵热力图
    需要 confusion_matrix 数据 (从 evaluation_metrics.json 或重新计算)
    """
    # 尝试从文件加载
    cm_path = os.path.join(result_dir, 'confusion_matrix.json')
    eval_path = os.path.join(result_dir, 'evaluation_metrics.json')

    cm = None
    if os.path.exists(cm_path):
        with open(cm_path, encoding='utf-8') as f:
            cm = np.array(json.load(f))
    elif os.path.exists(eval_path):
        with open(eval_path, encoding='utf-8') as f:
            eval_data = json.load(f)
            if 'confusion_matrix' in eval_data:
                cm = np.array(eval_data['confusion_matrix'])

    if cm is None:
        print("  [Fig3] 跳过: 无混淆矩阵数据")
        return

    num_classes = cm.shape[0]
    classes = summary.get('dataset_info', {}).get('classes', [f'C{i}' for i in range(num_classes)])

    # 归一化 (按行)
    cm_norm = cm.astype('float') / cm.sum(axis=1, keepdims=True)
    cm_norm = np.nan_to_num(cm_norm)

    fig, ax = plt.subplots(figsize=(4.5, 4.0))

    im = ax.imshow(cm_norm, cmap='Blues', vmin=0, vmax=1, aspect='equal')

    # 标注数值
    for i in range(num_classes):
        for j in range(num_classes):
            val = cm_norm[i, j]
            if val > 0.01:
                color = 'white' if val > 0.5 else 'black'
                ax.text(j, i, f'{val:.2f}', ha='center', va='center',
                        fontsize=6, color=color)

    ax.set_xlabel('Predicted Label')
    ax.set_ylabel('True Label')
    ax.set_xticks(range(num_classes))
    ax.set_yticks(range(num_classes))
    ax.set_xticklabels(classes, rotation=45, ha='right', fontsize=6)
    ax.set_yticklabels(classes, fontsize=6)

    cbar = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
    cbar.set_label('Accuracy', fontsize=9)

    fig.tight_layout()
    fig.savefig(save_path)
    plt.close(fig)
    print(f"  [Fig3] Confusion matrix -> {save_path}")


# ============================================================
# Fig 4: ROC 曲线
# ============================================================

def plot_roc_curves(result_dir, summary, save_path):
    """
    绘制多类别 One-vs-Rest ROC 曲线
    需要 roc_data 或 all_outputs + all_labels
    """
    roc_path = os.path.join(result_dir, 'roc_data.json')
    eval_path = os.path.join(result_dir, 'evaluation_metrics.json')

    roc_data = None
    if os.path.exists(roc_path):
        with open(roc_path, encoding='utf-8') as f:
            roc_data = json.load(f)
    elif os.path.exists(eval_path):
        with open(eval_path, encoding='utf-8') as f:
            eval_data = json.load(f)
            if 'class_auc' in eval_data:
                roc_data = eval_data

    if roc_data is None or 'class_auc' not in roc_data:
        print("  [Fig4] 跳过: 无 ROC 数据")
        return

    class_auc = roc_data['class_auc']
    macro_auc = roc_data.get('macro_auc', np.mean(list(class_auc.values())))
    num_classes = len(class_auc)
    classes = summary.get('dataset_info', {}).get('classes', [f'Class {i}' for i in range(num_classes)])

    # 如果有逐类 fpr/tpr 数据
    fig, ax = plt.subplots(figsize=(4.0, 3.5))

    if 'roc_curves' in roc_data:
        for i, (cls_name, curve) in enumerate(roc_data['roc_curves'].items()):
            color = FOLD_COLORS[i % len(FOLD_COLORS)]
            ax.plot(curve['fpr'], curve['tpr'], color=color, linewidth=0.8,
                    label=f'{classes[int(cls_name)]} (AUC={class_auc[cls_name]:.3f})')
    else:
        # 只画对角线和标注 macro AUC
        for i, (cls_idx, auc_val) in enumerate(class_auc.items()):
            color = FOLD_COLORS[i % len(FOLD_COLORS)]
            ax.plot([0, 1], [0, 1], color='gray', linestyle='--', linewidth=0.5)

    ax.plot([0, 1], [0, 1], 'k--', linewidth=0.5, alpha=0.5, label='Random (AUC=0.5)')
    ax.set_xlabel('False Positive Rate')
    ax.set_ylabel('True Positive Rate')
    ax.set_xlim(-0.02, 1.02)
    ax.set_ylim(-0.02, 1.02)
    ax.legend(loc='lower right', fontsize=7, ncol=1)

    ax.text(0.98, 0.60, f'Macro AUC = {macro_auc:.4f}',
            transform=ax.transAxes, ha='right', va='top',
            fontsize=9, fontweight='bold',
            bbox=dict(boxstyle='round,pad=0.3', facecolor='lightyellow',
                      edgecolor='gray', alpha=0.9))

    fig.tight_layout()
    fig.savefig(save_path)
    plt.close(fig)
    print(f"  [Fig4] ROC curves -> {save_path}")


# ============================================================
# Fig 5: t-SNE 嵌入可视化
# ============================================================

def plot_tsne_embedding(result_dir, summary, save_path):
    """
    绘制 t-SNE 2D 嵌入散点图
    需要 embedding_data.json (包含 embeddings + labels)
    """
    tsne_path = os.path.join(result_dir, 'tsne_data.json')
    emb_path = os.path.join(result_dir, 'embedding_data.json')

    data = None
    if os.path.exists(tsne_path):
        with open(tsne_path, encoding='utf-8') as f:
            data = json.load(f)
    elif os.path.exists(emb_path):
        # 需要做 t-SNE 降维
        with open(emb_path, encoding='utf-8') as f:
            raw = json.load(f)
        embeddings = np.array(raw['embeddings'])
        labels = np.array(raw['labels'])

        if embeddings.shape[1] > 2:
            from sklearn.manifold import TSNE
            tsne = TSNE(n_components=2, perplexity=30, random_state=42, n_iter=1000)
            embeddings_2d = tsne.fit_transform(embeddings)
        else:
            embeddings_2d = embeddings

        data = {
            'embeddings_2d': embeddings_2d.tolist(),
            'labels': labels.tolist()
        }

    if data is None:
        print("  [Fig5] 跳过: 无嵌入数据")
        return

    embeddings_2d = np.array(data['embeddings_2d'])
    labels = np.array(data['labels'])
    unique_labels = sorted(set(labels))
    classes = summary.get('dataset_info', {}).get('classes', [f'Class {i}' for i in unique_labels])

    fig, ax = plt.subplots(figsize=(4.0, 3.5))

    # 用不同颜色和标记绘制每个类别
    markers = ['o', 's', '^', 'D', 'v', 'p', 'h', '*', 'P', 'X',
               'd', '<', '>', '8', '1', '2', '3', '4', '+', 'x']

    for i, label in enumerate(unique_labels):
        mask = labels == label
        color = FOLD_COLORS[i % len(FOLD_COLORS)]
        marker = markers[i % len(markers)]
        cls_name = classes[i] if i < len(classes) else f'Class {label}'
        ax.scatter(embeddings_2d[mask, 0], embeddings_2d[mask, 1],
                   c=color, marker=marker, s=25, alpha=0.7, edgecolors='black',
                   linewidth=0.3, label=cls_name, zorder=3)

    ax.set_xlabel('t-SNE Dimension 1')
    ax.set_ylabel('t-SNE Dimension 2')
    ax.legend(loc='center left', bbox_to_anchor=(1.0, 0.5), fontsize=6, ncol=1,
              frameon=True, markerscale=1.2)

    fig.tight_layout()
    fig.savefig(save_path)
    plt.close(fig)
    print(f"  [Fig5] t-SNE embedding -> {save_path}")


# ============================================================
# Fig 6: 各类别精确率/召回率/F1 雷达图 (可选)
# ============================================================

def plot_per_class_metrics(result_dir, summary, save_path):
    """
    绘制每个类别的 Precision / Recall / F1 分组柱状图
    需要 classification_report 数据
    """
    eval_path = os.path.join(result_dir, 'evaluation_metrics.json')
    if not os.path.exists(eval_path):
        print("  [Fig6] 跳过: 无 evaluation_metrics.json")
        return

    with open(eval_path, encoding='utf-8') as f:
        eval_data = json.load(f)

    if 'classification_report' not in eval_data:
        print("  [Fig6] 跳过: 无 classification_report 数据")
        return

    report = eval_data['classification_report']
    classes = []
    precisions = []
    recalls = []
    f1s = []

    for cls_name, metrics in report.items():
        if cls_name.startswith('class_'):
            classes.append(cls_name)
            precisions.append(metrics.get('precision', 0))
            recalls.append(metrics.get('recall', 0))
            f1s.append(metrics.get('f1-score', 0))

    if not classes:
        return

    n = len(classes)
    x = np.arange(n)
    width = 0.25

    fig, ax = plt.subplots(figsize=(7.0, 3.0))

    bars1 = ax.bar(x - width, precisions, width, label='Precision',
                   color=COLORS['primary'], edgecolor='black', linewidth=0.5, alpha=0.85)
    bars2 = ax.bar(x, recalls, width, label='Recall',
                   color=COLORS['secondary'], edgecolor='black', linewidth=0.5, alpha=0.85)
    bars3 = ax.bar(x + width, f1s, width, label='F1-Score',
                   color=COLORS['tertiary'], edgecolor='black', linewidth=0.5, alpha=0.85)

    ax.set_xlabel('Class')
    ax.set_ylabel('Score')
    ax.set_xticks(x)
    ax.set_xticklabels(classes, rotation=45, ha='right', fontsize=7)
    ax.set_ylim(0, 1.1)
    ax.legend(loc='lower right', ncol=3)

    fig.tight_layout()
    fig.savefig(save_path)
    plt.close(fig)
    print(f"  [Fig6] Per-class metrics -> {save_path}")


# ============================================================
# 主函数
# ============================================================

def main():
    parser = argparse.ArgumentParser(description='SCI 论文级图表绘制')
    parser.add_argument('--result_dir', type=str, required=True,
                        help='训练结果目录 (包含 kfold_results_summary.json)')
    args = parser.parse_args()

    set_sci_style()

    result_dir = args.result_dir
    summary_path = os.path.join(result_dir, 'kfold_results_summary.json')

    if not os.path.exists(summary_path):
        # 也尝试 arcface
        summary_path = os.path.join(result_dir, 'arcface_summary.json')

    if not os.path.exists(summary_path):
        print(f"错误: 在 {result_dir} 中找不到结果文件")
        sys.exit(1)

    with open(summary_path, encoding='utf-8') as f:
        summary = json.load(f)

    print(f"加载结果: {summary_path}")
    print(f"输出目录: {result_dir}")
    print("=" * 50)

    # 按顺序生成各图
    plot_kfold_accuracy(result_dir, summary,
                        os.path.join(result_dir, 'fig1_kfold_accuracy.png'))

    plot_learning_curves(result_dir, summary,
                         os.path.join(result_dir, 'fig2_learning_curves.png'))

    plot_confusion_matrix(result_dir, summary,
                          os.path.join(result_dir, 'fig3_confusion_matrix.png'))

    plot_roc_curves(result_dir, summary,
                    os.path.join(result_dir, 'fig4_roc_curves.png'))

    plot_tsne_embedding(result_dir, summary,
                        os.path.join(result_dir, 'fig5_tsne_embedding.png'))

    plot_per_class_metrics(result_dir, summary,
                           os.path.join(result_dir, 'fig6_per_class_metrics.png'))

    print("=" * 50)
    print("绘图完成！")


if __name__ == "__main__":
    main()