|
|
#!/usr/bin/env python3
|
|
|
# -*- coding: utf-8 -*-
|
|
|
"""
|
|
|
训练过程可视化脚本
|
|
|
|
|
|
该脚本用于可视化One-Prompt模型的训练过程，包括：
|
|
|
1. 训练/验证损失曲线
|
|
|
2. IoU和Dice指标曲线
|
|
|
3. 学习率变化曲线
|
|
|
4. 分割结果可视化
|
|
|
|
|
|
Usage:
|
|
|
python scripts/visualize_training.py --log_dir logs/polyp_val_test_2025_12_16_23_52_30
|
|
|
"""
|
|
|
|
|
|
import os
|
|
|
import re
|
|
|
import argparse
|
|
|
import matplotlib.pyplot as plt
|
|
|
import numpy as np
|
|
|
from pathlib import Path
|
|
|
from typing import List, Tuple, Dict
|
|
|
import matplotlib
|
|
|
matplotlib.use('Agg') # 非GUI后端
|
|
|
|
|
|
# 设置中文字体
|
|
|
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
|
|
|
plt.rcParams['axes.unicode_minus'] = False
|
|
|
|
|
|
|
|
|
def parse_log_file(log_path: str) -> Dict[str, List[float]]:
|
|
|
"""
|
|
|
解析训练日志文件。
|
|
|
|
|
|
Args:
|
|
|
log_path: 日志文件路径
|
|
|
|
|
|
Returns:
|
|
|
包含训练指标的字典
|
|
|
"""
|
|
|
metrics = {
|
|
|
'epoch': [],
|
|
|
'train_loss': [],
|
|
|
'val_loss': [],
|
|
|
'iou': [],
|
|
|
'dice': []
|
|
|
}
|
|
|
|
|
|
with open(log_path, 'r', encoding='utf-8') as f:
|
|
|
for line in f:
|
|
|
# 解析训练损失: Train loss: 0.455222487449646|| @ epoch 0.
|
|
|
train_match = re.search(r'Train loss: ([\d.]+)\|\| @ epoch (\d+)', line)
|
|
|
if train_match:
|
|
|
loss = float(train_match.group(1))
|
|
|
epoch = int(train_match.group(2))
|
|
|
if epoch >= len(metrics['train_loss']):
|
|
|
metrics['train_loss'].append(loss)
|
|
|
metrics['epoch'].append(epoch)
|
|
|
|
|
|
# 解析验证指标: Total score: 0.367, IOU: 0.012, DICE: 0.022 || @ epoch 2.
|
|
|
val_match = re.search(
|
|
|
r'Total score: ([\d.]+), IOU: ([\d.]+), DICE: ([\d.]+) \|\| @ epoch (\d+)',
|
|
|
line
|
|
|
)
|
|
|
if val_match:
|
|
|
val_loss = float(val_match.group(1))
|
|
|
iou = float(val_match.group(2))
|
|
|
dice = float(val_match.group(3))
|
|
|
metrics['val_loss'].append(val_loss)
|
|
|
metrics['iou'].append(iou)
|
|
|
metrics['dice'].append(dice)
|
|
|
|
|
|
return metrics
|
|
|
|
|
|
|
|
|
def plot_loss_curves(metrics: Dict[str, List[float]], save_path: str) -> None:
|
|
|
"""
|
|
|
绘制训练和验证损失曲线。
|
|
|
|
|
|
Args:
|
|
|
metrics: 训练指标字典
|
|
|
save_path: 图像保存路径
|
|
|
"""
|
|
|
fig, ax = plt.subplots(figsize=(10, 6))
|
|
|
|
|
|
epochs = metrics['epoch']
|
|
|
train_loss = metrics['train_loss']
|
|
|
|
|
|
# 绘制训练损失
|
|
|
ax.plot(epochs, train_loss, 'b-', label='Training Loss', linewidth=2)
|
|
|
|
|
|
# 如果有验证损失，绘制验证损失
|
|
|
if metrics['val_loss']:
|
|
|
# 验证是每隔几个epoch进行的，需要对齐x轴
|
|
|
val_epochs = np.linspace(0, max(epochs), len(metrics['val_loss']))
|
|
|
ax.plot(val_epochs, metrics['val_loss'], 'r--', label='Validation Loss', linewidth=2)
|
|
|
|
|
|
ax.set_xlabel('Epoch', fontsize=12)
|
|
|
ax.set_ylabel('Loss', fontsize=12)
|
|
|
ax.set_title('Training and Validation Loss Curves', fontsize=14)
|
|
|
ax.legend(loc='upper right', fontsize=10)
|
|
|
ax.grid(True, alpha=0.3)
|
|
|
|
|
|
# 添加最佳损失标注
|
|
|
if train_loss:
|
|
|
min_idx = np.argmin(train_loss)
|
|
|
ax.annotate(f'Best: {train_loss[min_idx]:.4f}',
|
|
|
xy=(epochs[min_idx], train_loss[min_idx]),
|
|
|
xytext=(epochs[min_idx] + 2, train_loss[min_idx] + 0.1),
|
|
|
arrowprops=dict(arrowstyle='->', color='blue'),
|
|
|
fontsize=10, color='blue')
|
|
|
|
|
|
plt.tight_layout()
|
|
|
plt.savefig(save_path, dpi=150, bbox_inches='tight')
|
|
|
plt.close()
|
|
|
print(f"损失曲线已保存至: {save_path}")
|
|
|
|
|
|
|
|
|
def plot_metric_curves(metrics: Dict[str, List[float]], save_path: str) -> None:
|
|
|
"""
|
|
|
绘制IoU和Dice指标曲线。
|
|
|
|
|
|
Args:
|
|
|
metrics: 训练指标字典
|
|
|
save_path: 图像保存路径
|
|
|
"""
|
|
|
if not metrics['iou'] or not metrics['dice']:
|
|
|
print("警告: 没有IoU/Dice指标数据")
|
|
|
return
|
|
|
|
|
|
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))
|
|
|
|
|
|
val_epochs = range(len(metrics['iou']))
|
|
|
|
|
|
# IoU曲线
|
|
|
ax1.plot(val_epochs, metrics['iou'], 'g-o', label='IoU', linewidth=2, markersize=6)
|
|
|
ax1.set_xlabel('Validation Step', fontsize=12)
|
|
|
ax1.set_ylabel('IoU', fontsize=12)
|
|
|
ax1.set_title('Intersection over Union (IoU)', fontsize=14)
|
|
|
ax1.grid(True, alpha=0.3)
|
|
|
|
|
|
# 标注最佳IoU
|
|
|
if metrics['iou']:
|
|
|
max_idx = np.argmax(metrics['iou'])
|
|
|
max_iou = metrics['iou'][max_idx]
|
|
|
ax1.annotate(f'Best: {max_iou:.4f}',
|
|
|
xy=(max_idx, max_iou),
|
|
|
xytext=(max_idx + 0.5, max_iou + 0.01),
|
|
|
arrowprops=dict(arrowstyle='->', color='green'),
|
|
|
fontsize=10, color='green')
|
|
|
|
|
|
# Dice曲线
|
|
|
ax2.plot(val_epochs, metrics['dice'], 'm-s', label='Dice', linewidth=2, markersize=6)
|
|
|
ax2.set_xlabel('Validation Step', fontsize=12)
|
|
|
ax2.set_ylabel('Dice Score', fontsize=12)
|
|
|
ax2.set_title('Dice Coefficient', fontsize=14)
|
|
|
ax2.grid(True, alpha=0.3)
|
|
|
|
|
|
# 标注最佳Dice
|
|
|
if metrics['dice']:
|
|
|
max_idx = np.argmax(metrics['dice'])
|
|
|
max_dice = metrics['dice'][max_idx]
|
|
|
ax2.annotate(f'Best: {max_dice:.4f}',
|
|
|
xy=(max_idx, max_dice),
|
|
|
xytext=(max_idx + 0.5, max_dice + 0.01),
|
|
|
arrowprops=dict(arrowstyle='->', color='purple'),
|
|
|
fontsize=10, color='purple')
|
|
|
|
|
|
plt.tight_layout()
|
|
|
plt.savefig(save_path, dpi=150, bbox_inches='tight')
|
|
|
plt.close()
|
|
|
print(f"指标曲线已保存至: {save_path}")
|
|
|
|
|
|
|
|
|
def plot_combined_dashboard(metrics: Dict[str, List[float]], save_path: str) -> None:
|
|
|
"""
|
|
|
绘制综合训练仪表板。
|
|
|
|
|
|
Args:
|
|
|
metrics: 训练指标字典
|
|
|
save_path: 图像保存路径
|
|
|
"""
|
|
|
fig = plt.figure(figsize=(16, 10))
|
|
|
|
|
|
# 创建子图布局
|
|
|
gs = fig.add_gridspec(2, 3, hspace=0.3, wspace=0.3)
|
|
|
|
|
|
# 1. 训练损失曲线
|
|
|
ax1 = fig.add_subplot(gs[0, 0])
|
|
|
if metrics['train_loss']:
|
|
|
ax1.plot(metrics['epoch'], metrics['train_loss'], 'b-', linewidth=2)
|
|
|
ax1.fill_between(metrics['epoch'], metrics['train_loss'], alpha=0.3)
|
|
|
ax1.set_title('Training Loss', fontsize=12, fontweight='bold')
|
|
|
ax1.set_xlabel('Epoch')
|
|
|
ax1.set_ylabel('Loss')
|
|
|
ax1.grid(True, alpha=0.3)
|
|
|
|
|
|
# 2. 验证损失曲线
|
|
|
ax2 = fig.add_subplot(gs[0, 1])
|
|
|
if metrics['val_loss']:
|
|
|
ax2.plot(range(len(metrics['val_loss'])), metrics['val_loss'], 'r-', linewidth=2)
|
|
|
ax2.fill_between(range(len(metrics['val_loss'])), metrics['val_loss'], alpha=0.3, color='red')
|
|
|
ax2.set_title('Validation Loss', fontsize=12, fontweight='bold')
|
|
|
ax2.set_xlabel('Validation Step')
|
|
|
ax2.set_ylabel('Loss')
|
|
|
ax2.grid(True, alpha=0.3)
|
|
|
|
|
|
# 3. IoU曲线
|
|
|
ax3 = fig.add_subplot(gs[0, 2])
|
|
|
if metrics['iou']:
|
|
|
ax3.plot(range(len(metrics['iou'])), metrics['iou'], 'g-o', linewidth=2, markersize=4)
|
|
|
ax3.set_title('IoU Score', fontsize=12, fontweight='bold')
|
|
|
ax3.set_xlabel('Validation Step')
|
|
|
ax3.set_ylabel('IoU')
|
|
|
ax3.grid(True, alpha=0.3)
|
|
|
|
|
|
# 4. Dice曲线
|
|
|
ax4 = fig.add_subplot(gs[1, 0])
|
|
|
if metrics['dice']:
|
|
|
ax4.plot(range(len(metrics['dice'])), metrics['dice'], 'm-s', linewidth=2, markersize=4)
|
|
|
ax4.set_title('Dice Score', fontsize=12, fontweight='bold')
|
|
|
ax4.set_xlabel('Validation Step')
|
|
|
ax4.set_ylabel('Dice')
|
|
|
ax4.grid(True, alpha=0.3)
|
|
|
|
|
|
# 5. 损失对比柱状图
|
|
|
ax5 = fig.add_subplot(gs[1, 1])
|
|
|
if metrics['train_loss'] and metrics['val_loss']:
|
|
|
x = np.arange(2)
|
|
|
values = [np.mean(metrics['train_loss']), np.mean(metrics['val_loss'])]
|
|
|
bars = ax5.bar(x, values, color=['blue', 'red'], alpha=0.7)
|
|
|
ax5.set_xticks(x)
|
|
|
ax5.set_xticklabels(['Avg Train Loss', 'Avg Val Loss'])
|
|
|
ax5.set_title('Average Loss Comparison', fontsize=12, fontweight='bold')
|
|
|
# 添加数值标签
|
|
|
for bar, val in zip(bars, values):
|
|
|
ax5.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01,
|
|
|
f'{val:.4f}', ha='center', va='bottom', fontsize=10)
|
|
|
|
|
|
# 6. 训练统计信息
|
|
|
ax6 = fig.add_subplot(gs[1, 2])
|
|
|
ax6.axis('off')
|
|
|
|
|
|
# 计算统计信息
|
|
|
stats_text = "Training Statistics\n" + "="*30 + "\n\n"
|
|
|
if metrics['train_loss']:
|
|
|
stats_text += f"Total Epochs: {len(metrics['epoch'])}\n"
|
|
|
stats_text += f"Final Train Loss: {metrics['train_loss'][-1]:.4f}\n"
|
|
|
stats_text += f"Best Train Loss: {min(metrics['train_loss']):.4f}\n"
|
|
|
if metrics['val_loss']:
|
|
|
stats_text += f"Final Val Loss: {metrics['val_loss'][-1]:.4f}\n"
|
|
|
stats_text += f"Best Val Loss: {min(metrics['val_loss']):.4f}\n"
|
|
|
if metrics['iou']:
|
|
|
stats_text += f"Best IoU: {max(metrics['iou']):.4f}\n"
|
|
|
if metrics['dice']:
|
|
|
stats_text += f"Best Dice: {max(metrics['dice']):.4f}\n"
|
|
|
|
|
|
ax6.text(0.1, 0.5, stats_text, transform=ax6.transAxes, fontsize=11,
|
|
|
verticalalignment='center', fontfamily='monospace',
|
|
|
bbox=dict(boxstyle='round', facecolor='lightgray', alpha=0.5))
|
|
|
|
|
|
# 添加总标题
|
|
|
fig.suptitle('One-Prompt Medical Image Segmentation - Training Dashboard',
|
|
|
fontsize=16, fontweight='bold', y=0.98)
|
|
|
|
|
|
plt.savefig(save_path, dpi=150, bbox_inches='tight')
|
|
|
plt.close()
|
|
|
print(f"训练仪表板已保存至: {save_path}")
|
|
|
|
|
|
|
|
|
def main():
|
|
|
"""主函数。"""
|
|
|
parser = argparse.ArgumentParser(description='可视化训练过程')
|
|
|
parser.add_argument('--log_dir', type=str, required=True, help='日志目录路径')
|
|
|
parser.add_argument('--output_dir', type=str, default=None, help='输出目录')
|
|
|
args = parser.parse_args()
|
|
|
|
|
|
log_dir = Path(args.log_dir)
|
|
|
output_dir = Path(args.output_dir) if args.output_dir else log_dir / 'visualizations'
|
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
# 查找日志文件
|
|
|
log_files = list(log_dir.glob('Log/*.log'))
|
|
|
if not log_files:
|
|
|
print(f"错误: 在 {log_dir}/Log/ 目录下未找到日志文件")
|
|
|
return
|
|
|
|
|
|
log_path = log_files[0]
|
|
|
print(f"正在解析日志文件: {log_path}")
|
|
|
|
|
|
# 解析日志
|
|
|
metrics = parse_log_file(str(log_path))
|
|
|
print(f"解析完成: {len(metrics['epoch'])} 个epoch, "
|
|
|
f"{len(metrics['val_loss'])} 次验证")
|
|
|
|
|
|
# 生成可视化
|
|
|
plot_loss_curves(metrics, str(output_dir / 'loss_curves.png'))
|
|
|
plot_metric_curves(metrics, str(output_dir / 'metric_curves.png'))
|
|
|
plot_combined_dashboard(metrics, str(output_dir / 'training_dashboard.png'))
|
|
|
|
|
|
print(f"\n所有可视化已保存至: {output_dir}")
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
main()
|
