564 lines
27 KiB
Python
564 lines
27 KiB
Python
import logging, argparse, sys, utils
|
||
import shutil, tempfile
|
||
from pathlib import Path
|
||
from typing import Dict, List, Tuple
|
||
|
||
import albumentations as A
|
||
import cv2
|
||
import matplotlib.pyplot as plt
|
||
import numpy as np
|
||
import segmentation_models_pytorch as smp
|
||
import torch
|
||
from albumentations.pytorch import ToTensorV2
|
||
from scipy.special import softmax
|
||
from sklearn.metrics import (auc, average_precision_score,
|
||
precision_recall_curve, roc_curve)
|
||
from tqdm import tqdm
|
||
|
||
# 本地应用/库的导入
|
||
import config
|
||
from utils import log_metrics_to_csv, get_preprocessing_transform, setup_directories
|
||
|
||
# --- 日志设置 ---
|
||
# 配置日志记录器
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format="%(asctime)s [%(levelname)s] %(message)s",
|
||
handlers=[logging.StreamHandler()] # 输出日志到控制台
|
||
)
|
||
# 为 matplotlib 设置中文字体，请根据您的系统环境选择合适的字体
|
||
plt.rcParams['font.sans-serif'] = ['Noto Sans CJK JP', 'SimHei', 'Microsoft YaHei'] # 例如：SimHei, Microsoft YaHei
|
||
plt.rcParams['axes.unicode_minus'] = False # 解决负号显示问题
|
||
|
||
# --- 核心工具函数 ---
|
||
|
||
# 将单通道的类别索引掩码转换为 RGB 彩色掩码
|
||
def colorize_mask(mask: np.ndarray, class_rgb_values: List[Tuple[int, int, int]]) -> np.ndarray:
|
||
"""将单通道的类别索引掩码转换为 RGB 彩色掩码。"""
|
||
color_mask = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)
|
||
# 一个通用的调色板，可以轻松扩展
|
||
palette = [
|
||
(0, 0, 0), (220, 20, 60), (0, 255, 0), (0, 0, 255),
|
||
(255, 255, 0), (255, 0, 255), (0, 255, 255), (244, 164, 96),
|
||
]
|
||
for class_idx, color in enumerate(palette):
|
||
if class_idx < len(class_rgb_values):
|
||
# 假设掩码中的像素值 [0, 1, 2...] 对应类别的索引
|
||
pixels_to_color = (mask == class_idx)
|
||
color_mask[pixels_to_color] = color
|
||
return color_mask
|
||
|
||
# --- 分析与可视化函数 ---
|
||
|
||
def generate_and_save_curves(y_true: np.ndarray, y_probs: np.ndarray, class_names: List[str], output_dir: Path) -> None:
|
||
"""
|
||
根据整个测试集的聚合预测结果，为每个类别生成并保存 ROC 和 PR 曲线。
|
||
"""
|
||
logging.info("正在生成 ROC 和 PR 曲线...")
|
||
num_classes = len(class_names)
|
||
plt.style.use('seaborn-v0_8-whitegrid')
|
||
|
||
# --- 准备画布 ---
|
||
fig_roc, ax_roc = plt.subplots(figsize=(10, 8))
|
||
ax_roc.plot([0, 1], [0, 1], 'k--', label='随机猜测')
|
||
fig_pr, ax_pr = plt.subplots(figsize=(10, 8))
|
||
|
||
# --- 为每个类别（跳过背景）生成曲线 ---
|
||
for i in range(1, num_classes):
|
||
class_name = class_names[i]
|
||
y_true_binary = (y_true == i).astype(int)
|
||
y_class_probs = y_probs[:, i]
|
||
|
||
# ROC 曲线
|
||
fpr, tpr, _ = roc_curve(y_true_binary, y_class_probs)
|
||
roc_auc = auc(fpr, tpr)
|
||
ax_roc.plot(fpr, tpr, label=f'类别 {class_name} (AUC = {roc_auc:.4f})')
|
||
|
||
# PR 曲线
|
||
precision, recall, _ = precision_recall_curve(y_true_binary, y_class_probs)
|
||
avg_precision = average_precision_score(y_true_binary, y_class_probs)
|
||
ax_pr.plot(recall, precision, label=f'类别 {class_name} (AP = {avg_precision:.4f})')
|
||
|
||
# --- 美化并保存 ROC 曲线图 ---
|
||
ax_roc.set(xlabel='False positive rate (FPR)', ylabel='True positive rate (TPR)', title='Receiver operating characteristic (ROC) curve')
|
||
ax_roc.legend(loc='lower right')
|
||
ax_roc.set_aspect('equal', adjustable='box')
|
||
fig_roc.tight_layout()
|
||
roc_save_path = output_dir / "ROC_Curves.png"
|
||
fig_roc.savefig(roc_save_path)
|
||
plt.close(fig_roc)
|
||
logging.info(f"-> ROC 曲线图已保存至: {roc_save_path}")
|
||
|
||
# --- 美化并保存 PR 曲线图 ---
|
||
ax_pr.set(xlabel='Recall', ylabel='Precision', title='Precision-Recall (PR) Curve')
|
||
ax_pr.legend(loc='lower left')
|
||
ax_pr.set_aspect('equal', adjustable='box')
|
||
fig_pr.tight_layout()
|
||
pr_save_path = output_dir / "Precision_Recall_Curves.png"
|
||
fig_pr.savefig(pr_save_path)
|
||
plt.close(fig_pr)
|
||
logging.info(f"-> PR 曲线图已保存至: {pr_save_path}")
|
||
|
||
def save_visual_comparison(
|
||
image_name: str, original_image: np.ndarray, pred_mask: np.ndarray,
|
||
gt_mask: np.ndarray = None, stats: Dict = None, stats_text: str = "",
|
||
save_dir: str = ""
|
||
) -> None:
|
||
"""保存原始图像、预测掩码以及可选的真实掩码的对比图。"""
|
||
pred_color_mask = colorize_mask(pred_mask, config.CLASS_RGB_VALUES)
|
||
|
||
if gt_mask is not None and stats is not None:
|
||
num_plots = 3
|
||
figsize = (22, 8)
|
||
gt_color_mask = colorize_mask(gt_mask, config.CLASS_RGB_VALUES)
|
||
save_name = f"{Path(image_name).stem}-IOU_{stats['iou']:.4f}-ACC_{stats['acc']:.4f}.png"
|
||
else:
|
||
num_plots = 2
|
||
figsize = (14, 6)
|
||
save_name = f"{Path(image_name).stem}.png"
|
||
|
||
fig, axes = plt.subplots(1, num_plots, figsize=figsize)
|
||
fig.suptitle(f"Predict: {image_name}", fontsize=16)
|
||
|
||
axes[0].imshow(original_image)
|
||
axes[0].set_title("ORI_IMG")
|
||
axes[0].axis('off')
|
||
|
||
pred_title = "Predicted Mask"
|
||
if stats:
|
||
pred_title += f"\nIoU: {stats['iou']:.4f} | Acc: {stats['acc']:.4f}"
|
||
axes[1].imshow(pred_color_mask)
|
||
axes[1].set_title(pred_title)
|
||
axes[1].axis('off')
|
||
|
||
if num_plots == 3:
|
||
axes[2].imshow(gt_color_mask)
|
||
axes[2].set_title("Ground Truth")
|
||
axes[2].axis('off')
|
||
plt.figtext(0.5, 0.01, stats_text, ha="center", fontsize=10,
|
||
bbox={"facecolor":"white", "alpha":0.7, "pad":5}, family='monospace')
|
||
fig.subplots_adjust(bottom=0.3)
|
||
|
||
plt.tight_layout()
|
||
plt.savefig(save_dir / save_name)
|
||
plt.close(fig)
|
||
|
||
# --- 主要逻辑函数 ---
|
||
|
||
def load_model(model_architecture: str, model_save_path: Path) -> torch.nn.Module:
|
||
"""
|
||
加载已训练的分割模型。
|
||
此函数经过优化，可以智能处理由 torch.nn.DataParallel 保存的
|
||
(带有 'module.' 前缀) 和常规保存的模型权重。
|
||
"""
|
||
logging.info(f"正在从以下路径加载模型: {model_save_path}")
|
||
num_classes = len(config.CLASSES)
|
||
|
||
# 1. 【修改】根据 config 动态创建基础模型架构
|
||
logging.info(f"正在重建模型架构: '{model_architecture}'")
|
||
try:
|
||
model_params = config.ALL_MODEL_CONFIGS[model_architecture].copy() #
|
||
model_class = getattr(smp, model_architecture)
|
||
except AttributeError:
|
||
logging.error(f"模型 '{model_architecture}' 在 segmentation_models_pytorch 库中不存在！")
|
||
raise
|
||
|
||
# 2. 【修改】准备参数，与训练时保持一致
|
||
# 注意：加载权重时，我们不希望再次下载预训练权重，所以设为 None
|
||
params = model_params
|
||
params['in_channels'] = 3
|
||
params['classes'] = num_classes
|
||
params['encoder_weights'] = None # 非常重要！避免在预测时重新下载ImageNet权重
|
||
|
||
# ======================== 【新增代码段开始】 ======================== #
|
||
# 自动检测 encoder_name 是否包含 'vit'
|
||
encoder_name = model_params.get('encoder_name', '')
|
||
if 'vit' in encoder_name.lower():
|
||
params['dynamic_img_size'] = True
|
||
logging.info(f"检测到 ViT 编码器 ('{encoder_name}')。自动设置 dynamic_img_size=True。")
|
||
# ======================== 【新增代码段结束】 ======================== #
|
||
|
||
# 3. 【修改】使用 **kwargs 创建模型实例
|
||
model = model_class(**params).to(config.DEVICE)
|
||
|
||
# 4. 从文件加载 state_dict
|
||
state_dict = torch.load(model_save_path, map_location=torch.device(config.DEVICE))
|
||
|
||
# 3. 检查是否存在 'module.' 前缀 (判断是否为 DataParallel 保存的权重)
|
||
# 我们通过检查字典中任何一个键是否以 'module.' 开头来判断
|
||
is_dataparallel = any(key.startswith('module.') for key in state_dict.keys())
|
||
|
||
if is_dataparallel:
|
||
logging.info("检测到模型由 DataParallel 保存，将移除 'module.' 前缀。")
|
||
# 创建一个新的、不带前缀的 state_dict
|
||
new_state_dict = {key.replace('module.', ''): value for key, value in state_dict.items()}
|
||
model.load_state_dict(new_state_dict)
|
||
else:
|
||
# 如果没有 'module.' 前缀，直接加载
|
||
logging.info("直接加载标准模型权重。")
|
||
model.load_state_dict(state_dict)
|
||
|
||
model.eval()
|
||
return model
|
||
|
||
def calculate_and_save_final_metrics(stats_list: List[Dict], num_classes: int, save_dir: str) -> None:
|
||
"""聚合每张图像的统计数据，计算总体指标，并保存到 CSV 文件。"""
|
||
if not stats_list:
|
||
logging.warning("未收集到统计数据，跳过最终指标计算。")
|
||
return
|
||
|
||
logging.info("\n正在聚合结果以计算最终指标...")
|
||
# 堆叠所有统计张量并按批次维度求和
|
||
tp = torch.cat([s['tp'] for s in stats_list], dim=0).sum(dim=0) # (N×B, C) -> (C,)
|
||
fp = torch.cat([s['fp'] for s in stats_list], dim=0).sum(dim=0) # (N×B, C) -> (C,)
|
||
fn = torch.cat([s['fn'] for s in stats_list], dim=0).sum(dim=0) # (N×B, C) -> (C,)
|
||
tn = torch.cat([s['tn'] for s in stats_list], dim=0).sum(dim=0) # (N×B, C) -> (C,)
|
||
|
||
# ======================== 【新增代码段开始】 ======================== #
|
||
# 根据 config.py 中的设置，筛选出用于计算宏观指标的类别
|
||
ignore_indices = [config.CLASSES.index(cls) for cls in config.EVALUATION_CLASSES_TO_IGNORE if cls in config.CLASSES]
|
||
|
||
keep_mask = torch.ones(num_classes, dtype=torch.bool, device=tp.device)
|
||
if ignore_indices:
|
||
keep_mask[ignore_indices] = False
|
||
logging.info(f"最终评估时将忽略类别: {config.EVALUATION_CLASSES_TO_IGNORE}")
|
||
|
||
tp_filtered = tp[keep_mask]
|
||
fp_filtered = fp[keep_mask]
|
||
fn_filtered = fn[keep_mask]
|
||
tn_filtered = tn[keep_mask]
|
||
# ======================== 【新增代码段结束】 ======================== #
|
||
|
||
|
||
# ======================== 【修改代码段开始】 ======================== #
|
||
# 使用过滤后的统计数据计算总体指标
|
||
metrics = {
|
||
"iou_score": smp.metrics.iou_score(tp_filtered, fp_filtered, fn_filtered, tn_filtered, reduction='micro').item(),
|
||
"f1_score": smp.metrics.f1_score(tp_filtered, fp_filtered, fn_filtered, tn_filtered, reduction='micro').item(),
|
||
"accuracy": smp.metrics.accuracy(tp_filtered, fp_filtered, fn_filtered, tn_filtered, reduction='micro').item(),
|
||
"recall": smp.metrics.recall(tp_filtered, fp_filtered, fn_filtered, tn_filtered, reduction='micro').item(),
|
||
"precision": smp.metrics.precision(tp_filtered, fp_filtered, fn_filtered, tn_filtered, reduction='micro').item(),
|
||
}
|
||
# ======================== 【修改代码段结束】 ======================== #
|
||
|
||
# 添加每个类别的原始统计数据（使用未经过滤的数据）
|
||
for i in range(num_classes):
|
||
metrics[f'tp_class_{i}'] = tp[i].item()
|
||
metrics[f'tn_class_{i}'] = tn[i].item()
|
||
metrics[f'fp_class_{i}'] = fp[i].item()
|
||
metrics[f'fn_class_{i}'] = fn[i].item()
|
||
|
||
csv_save_path = save_dir / "test_set_metrics.csv"
|
||
log_metrics_to_csv(metrics, str(csv_save_path))
|
||
logging.info(f"-> 整体测试集指标已保存至: {csv_save_path}")
|
||
|
||
# 2. 【新增】一个帮助函数，用于交互式选择模型目录
|
||
def select_trained_model_path(base_dir: Path, model_architecture: str) -> Path:
|
||
"""
|
||
查找指定架构的所有训练运行目录，并让用户选择一个。
|
||
"""
|
||
logging.info(f"正在 '{base_dir}' 中搜索模型 '{model_architecture}' 的训练记录...")
|
||
|
||
# 查找所有以架构名开头的文件夹
|
||
run_dirs = sorted([d for d in base_dir.iterdir() if d.is_dir() and d.name.startswith(model_architecture+'_')])
|
||
|
||
if not run_dirs:
|
||
logging.error(f"未找到任何 '{model_architecture}' 模型的训练记录。请先运行 train.py。")
|
||
return None
|
||
|
||
print("\n请选择要用于预测的模型:")
|
||
for i, dir_path in enumerate(run_dirs):
|
||
print(f" [{i+1}] {dir_path.name}")
|
||
|
||
while True:
|
||
try:
|
||
choice = input(f"请输入选项编号 (1-{len(run_dirs)}) 或按 Enter 退出: ")
|
||
if not choice:
|
||
return None
|
||
choice_idx = int(choice) - 1
|
||
if 0 <= choice_idx < len(run_dirs):
|
||
selected_dir = run_dirs[choice_idx]
|
||
logging.info(f"已选择模型: {selected_dir}")
|
||
return selected_dir
|
||
else:
|
||
print("无效的选项，请重新输入。")
|
||
except (ValueError, IndexError):
|
||
print("无效的输入，请输入一个数字。")
|
||
|
||
# 【修改】函数签名和内部逻辑，不再接收一个列表，而是直接接收聚合后的统计数据
|
||
def calculate_and_save_final_metrics(tp, fp, fn, tn, num_classes: int, save_dir: Path) -> None:
|
||
"""根据聚合后的统计数据，计算总体指标，并保存到 CSV 文件。"""
|
||
logging.info("\n正在计算最终的整体指标...")
|
||
|
||
# 根据 config.py 中的设置，筛选出用于计算宏观指标的类别
|
||
ignore_indices = [config.CLASSES.index(cls) for cls in config.EVALUATION_CLASSES_TO_IGNORE if cls in config.CLASSES]
|
||
|
||
keep_mask = torch.ones(num_classes, dtype=torch.bool, device=tp.device)
|
||
if ignore_indices:
|
||
keep_mask[ignore_indices] = False
|
||
logging.info(f"最终评估时将忽略类别: {config.EVALUATION_CLASSES_TO_IGNORE}")
|
||
|
||
tp_filtered = tp[keep_mask]
|
||
fp_filtered = fp[keep_mask]
|
||
fn_filtered = fn[keep_mask]
|
||
tn_filtered = tn[keep_mask]
|
||
|
||
# 使用过滤后的统计数据计算总体指标
|
||
metrics = {
|
||
"iou_score": smp.metrics.iou_score(tp_filtered, fp_filtered, fn_filtered, tn_filtered, reduction='micro').item(),
|
||
"f1_score": smp.metrics.f1_score(tp_filtered, fp_filtered, fn_filtered, tn_filtered, reduction='micro').item(),
|
||
"accuracy": smp.metrics.accuracy(tp_filtered, fp_filtered, fn_filtered, tn_filtered, reduction='micro').item(),
|
||
"recall": smp.metrics.recall(tp_filtered, fp_filtered, fn_filtered, tn_filtered, reduction='micro').item(),
|
||
"precision": smp.metrics.precision(tp_filtered, fp_filtered, fn_filtered, tn_filtered, reduction='micro').item(),
|
||
}
|
||
|
||
# 添加每个类别的原始统计数据（使用未经过滤的数据）
|
||
for i in range(num_classes):
|
||
metrics[f'tp_class_{i}'] = tp[i].item()
|
||
metrics[f'tn_class_{i}'] = tn[i].item()
|
||
metrics[f'fp_class_{i}'] = fp[i].item()
|
||
metrics[f'fn_class_{i}'] = fn[i].item()
|
||
|
||
csv_save_path = save_dir / "test_set_metrics.csv"
|
||
log_metrics_to_csv(metrics, str(csv_save_path))
|
||
logging.info(f"-> 整体测试集指标已保存至: {csv_save_path}")
|
||
|
||
def main(model_architecture: str) -> None:
|
||
"""
|
||
主函数，执行模型加载、图像预测、评估，并生成分析结果。
|
||
"""
|
||
# --- 1. 【修改】选择模型并动态设置路径 ---
|
||
selected_run_dir = select_trained_model_path(Path(config.PREDICT_BEST_MODEL_DIR), model_architecture)
|
||
|
||
if not selected_run_dir:
|
||
logging.info("未选择任何模型，程序退出。")
|
||
sys.exit()
|
||
|
||
# 【重要】动态覆写 config 中的路径，使其指向所选模型的目录
|
||
# 这样，后续所有代码都会自动在正确的文件夹中加载模型和保存结果
|
||
best_model_save_path = selected_run_dir / "best_model.pth"
|
||
raw_mask_dir = selected_run_dir / "predicted_raw_masks"
|
||
analysis_results_dir = selected_run_dir / "prediction_analysis"
|
||
|
||
if not best_model_save_path.exists():
|
||
logging.error(f"错误: 在 '{selected_run_dir}' 中未找到 'best_model.pth' 文件！")
|
||
sys.exit()
|
||
|
||
# setup_directories 函数现在用于创建预测结果的子目录
|
||
utils.setup_directories(raw_mask_dir)
|
||
utils.setup_directories(analysis_results_dir)
|
||
|
||
model = load_model(model_architecture, best_model_save_path) # 1. 加载模型
|
||
|
||
# ======================== 【新增/修改的代码段开始】 ======================== #
|
||
# --- 2. 探测模型以确定正确的图像尺寸 ---
|
||
logging.info("正在探测模型以确定预测时所需的输入尺寸...") #
|
||
target_height, target_width = config.IMAGE_HEIGHT, config.IMAGE_WIDTH #
|
||
|
||
try:
|
||
# 检查编码器是否需要固定输入尺寸
|
||
if model.encoder.is_fixed_input_size:
|
||
required_size = model.encoder.input_size
|
||
target_height = required_size[1]
|
||
target_width = required_size[2]
|
||
logging.warning(
|
||
f"模型 '{model_architecture}' 的编码器需要固定的输入尺寸 "
|
||
f"({target_height}x{target_width})。"
|
||
f"将忽略 config.py 中的尺寸设置进行预测。"
|
||
)
|
||
else:
|
||
logging.info(f"模型 '{model_architecture}' 支持动态输入尺寸。将使用 config.py 中定义的尺寸 ({target_height}x{target_width})。")
|
||
except AttributeError:
|
||
logging.info("无法自动检测模型尺寸要求。将使用 config.py 中定义的尺寸。")
|
||
|
||
# --- 3. 获取用于预测的图像预处理流程 ---
|
||
# 使用探测到的或配置中指定的目标尺寸
|
||
transform = get_preprocessing_transform(target_height, target_width) #
|
||
# ======================== 【新增/修改的代码段结束】 ======================== #
|
||
|
||
# --- 4. 获取图像列表并判断模式 ---
|
||
image_filenames = sorted(list(config.TEST_IMAGE_DIR.iterdir()))
|
||
evaluate_mode = config.TEST_MASK_DIR.is_dir()
|
||
|
||
if evaluate_mode:
|
||
logging.info("正在以 [评估模式] 运行。将使用真实掩码进行评估。")
|
||
else:
|
||
logging.info("正在以 [仅预测模式] 运行。未找到真实掩码。")
|
||
|
||
# --- 3. 初始化用于增量计算的变量 ---
|
||
# 【修改】不再使用 list 来累积，而是直接累加 TP, FP, FN, TN
|
||
num_classes = len(config.CLASSES)
|
||
total_tp = torch.zeros(num_classes, dtype=torch.long, device=config.DEVICE)
|
||
total_fp = torch.zeros(num_classes, dtype=torch.long, device=config.DEVICE)
|
||
total_fn = torch.zeros(num_classes, dtype=torch.long, device=config.DEVICE)
|
||
total_tn = torch.zeros(num_classes, dtype=torch.long, device=config.DEVICE)
|
||
|
||
# 【新增】为 ROC/PR 曲线创建临时文件，用磁盘空间换内存
|
||
# 使用 tempfile 模块来安全地创建临时文件
|
||
temp_gt_file = tempfile.NamedTemporaryFile(delete=False, suffix='.npy')
|
||
temp_probs_file = tempfile.NamedTemporaryFile(delete=False, suffix='.npy')
|
||
logging.info(f"为ROC/PR曲线创建临时文件: {temp_gt_file.name}, {temp_probs_file.name}")
|
||
|
||
# --- 4. 主推理循环 - 预测图片 ---
|
||
for image_path in tqdm(image_filenames, desc="正在处理测试图像"):
|
||
# --- 4.1. 读取图片并预处理 ---
|
||
original_image = cv2.cvtColor(cv2.imread(str(image_path)), cv2.COLOR_BGR2RGB)
|
||
original_h, original_w = original_image.shape[:2]
|
||
image_tensor = transform(image=original_image)['image'].unsqueeze(0).to(config.DEVICE)
|
||
|
||
# --- 4.2. 图片推理并保存 ---
|
||
with torch.no_grad():
|
||
pred_logits_tensor = model(image_tensor)
|
||
pred_mask_tensor = torch.argmax(pred_logits_tensor, dim=1)
|
||
pred_mask_numpy = pred_mask_tensor.squeeze().cpu().numpy().astype(np.uint8)
|
||
pred_mask_resized_raw = cv2.resize(pred_mask_numpy, (original_w, original_h), interpolation=cv2.INTER_NEAREST)
|
||
|
||
# 保存原始（单通道）的预测掩码
|
||
cv2.imwrite(str(raw_mask_dir / f"{image_path.stem}.png"), pred_mask_resized_raw)
|
||
|
||
# --- 4.3. 评估与可视化 ---
|
||
gt_mask_numpy, per_image_stats, stats_text_display, gt_mask_raw = None, None, "", None
|
||
if evaluate_mode:
|
||
mask_path = config.TEST_MASK_DIR / image_path.name
|
||
if mask_path.exists():
|
||
gt_mask_raw = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE)
|
||
gt_mask_numpy = cv2.resize(
|
||
gt_mask_raw, (target_width, target_height),
|
||
interpolation=cv2.INTER_NEAREST
|
||
)
|
||
gt_mask_tensor = torch.from_numpy(gt_mask_numpy).long().to(config.DEVICE).unsqueeze(0)
|
||
|
||
# ======================== 【修复代码段开始】 ======================== #
|
||
|
||
# --- 1. [修复] 计算单张图像的统计数据 ---
|
||
# 这个代码块是必需的，它为下面的 per_image_stats 和 stats_text_display 定义了 tp, fp, fn, tn
|
||
tp, fp, fn, tn = smp.metrics.get_stats(
|
||
pred_mask_tensor, gt_mask_tensor,
|
||
mode=config.SEG_MODE, num_classes=num_classes
|
||
)
|
||
|
||
# --- 2. 累加到总数，用于最终的整体评估 ---
|
||
total_tp += tp.squeeze(0).to(config.DEVICE)
|
||
total_fp += fp.squeeze(0).to(config.DEVICE)
|
||
total_fn += fn.squeeze(0).to(config.DEVICE)
|
||
total_tn += tn.squeeze(0).to(config.DEVICE)
|
||
|
||
# --- 3. 将曲线数据写入临时文件 (这部分逻辑正确，无需修改) ---
|
||
probs_tensor = torch.softmax(pred_logits_tensor, dim=1)
|
||
gt_flat = gt_mask_tensor.cpu().numpy().flatten()
|
||
probs_flat = probs_tensor.permute(0, 2, 3, 1).reshape(-1, num_classes).cpu().numpy()
|
||
with open(temp_gt_file.name, 'ab') as f_gt, open(temp_probs_file.name, 'ab') as f_probs:
|
||
np.save(f_gt, gt_flat)
|
||
np.save(f_probs, probs_flat)
|
||
|
||
# --- 4. [修复] 现在可以安全地计算单图指标用于可视化 ---
|
||
# 因为上面的代码块已经定义了 tp, fp, fn, tn，所以这里不再报错
|
||
per_image_stats = {
|
||
'iou': smp.metrics.iou_score(tp, fp, fn, tn, reduction='micro').item(),
|
||
'acc': smp.metrics.accuracy(tp, fp, fn, tn, reduction='micro').item()
|
||
}
|
||
|
||
# --- 5. [修复] 修正生成统计文本时的 Bug ---
|
||
# `get_stats` 返回的 tp 张量形状是 (1, num_classes),
|
||
# 因此我们应该遍历第1维度 (类别)，并使用 [0, i] 进行索引
|
||
stats_text_display = "Each Class TP / FP / FN / TN:\n" + "-"*35 + "\n"
|
||
for i in range(1, tp.shape[1]): # 修正: 遍历类别维度 tp.shape[1]
|
||
if i < len(config.CLASSES):
|
||
class_name = config.CLASSES[i]
|
||
stats_text_display += (
|
||
f" {class_name:<10}: "
|
||
# 修正: 使用正确的索引 tp[0, i]
|
||
f"{tp[0, i]:>5} / {fp[0, i]:>5} / {fn[0, i]:>5} / {tn[0, i]:>7}\n"
|
||
)
|
||
|
||
save_visual_comparison(
|
||
image_name=image_path.name,
|
||
original_image=original_image,
|
||
pred_mask=pred_mask_resized_raw,
|
||
gt_mask=gt_mask_raw,
|
||
stats=per_image_stats,
|
||
stats_text=stats_text_display,
|
||
save_dir = analysis_results_dir
|
||
)
|
||
|
||
logging.info("\n预测流程已完成！")
|
||
logging.info(f"原始预测掩码已保存至: {raw_mask_dir}")
|
||
logging.info(f"分析图像已保存至: {analysis_results_dir}")
|
||
|
||
# --- 5. 最终分析（循环结束后） ---
|
||
if evaluate_mode:
|
||
# calculate_and_save_final_metrics(all_stats_for_metrics, len(config.CLASSES), save_dir = analysis_results_dir)
|
||
calculate_and_save_final_metrics(total_tp.cpu(), total_fp.cpu(), total_fn.cpu(), total_tn.cpu(), num_classes, save_dir=analysis_results_dir)
|
||
|
||
# if all_gt_for_curves:
|
||
# logging.info("\n正在聚合结果以生成 ROC 和 PR 曲线...")
|
||
# full_gt = torch.cat(all_gt_for_curves).numpy().flatten()
|
||
# full_logits = torch.cat(all_logits_for_curves)
|
||
|
||
# # 调整维度并计算概率
|
||
# num_classes = len(config.CLASSES)
|
||
# full_probs = softmax(
|
||
# full_logits.permute(0, 2, 3, 1).reshape(-1, num_classes).numpy(),
|
||
# axis=1
|
||
# )
|
||
|
||
# generate_and_save_curves(
|
||
# y_true=full_gt, y_probs=full_probs,
|
||
# class_names=config.CLASSES, output_dir=analysis_results_dir
|
||
# )
|
||
|
||
# 【新增】从临时文件中分块读取数据来生成曲线
|
||
try:
|
||
logging.info("\n正在从临时文件加载数据以生成 ROC 和 PR 曲线...")
|
||
temp_gt_file.close()
|
||
temp_probs_file.close()
|
||
|
||
all_gt_parts, all_probs_parts = [], []
|
||
with open(temp_gt_file.name, 'rb') as f_gt, open(temp_probs_file.name, 'rb') as f_probs:
|
||
while True:
|
||
try:
|
||
all_gt_parts.append(np.load(f_gt))
|
||
all_probs_parts.append(np.load(f_probs))
|
||
# 【修复】将 EOFError 添加到要捕获的异常列表中
|
||
except (IOError, ValueError, EOFError):
|
||
break # 文件读取完毕
|
||
|
||
if all_gt_parts:
|
||
full_gt = np.concatenate(all_gt_parts)
|
||
full_probs = np.concatenate(all_probs_parts)
|
||
|
||
generate_and_save_curves(
|
||
y_true=full_gt, y_probs=full_probs,
|
||
class_names=config.CLASSES, output_dir=analysis_results_dir
|
||
)
|
||
else:
|
||
logging.warning("临时文件中没有数据，跳过 ROC/PR 曲线生成。")
|
||
|
||
finally:
|
||
# 【新增】确保在最后删除临时文件
|
||
import os
|
||
os.remove(temp_gt_file.name)
|
||
os.remove(temp_probs_file.name)
|
||
logging.info("已清理临时文件。")
|
||
|
||
if __name__ == "__main__":
|
||
# 创建一个参数解析器
|
||
parser = argparse.ArgumentParser(description="训练图像分割模型。")
|
||
|
||
# 添加 --architecture 参数
|
||
# - `choices` 会自动从你的 config 文件中获取所有可用的模型名称
|
||
# - `default` 设置了在不提供参数时的默认模型
|
||
parser.add_argument(
|
||
"-a", "--architecture",
|
||
type=str,
|
||
# default='Unet',
|
||
choices=list(config.ALL_MODEL_CONFIGS.keys()),
|
||
required=True,
|
||
help="选择要训练的模型架构。"
|
||
)
|
||
|
||
# 解析命令行传入的参数
|
||
args = parser.parse_args()
|
||
|
||
# 将解析出的架构名称传入 main 函数并执行
|
||
main(model_architecture=args.architecture) |