""" 模型推理测试脚本 功能说明： 1. 加载训练好的SegFormer模型 2. 对单张图像进行分割预测 3. 可视化分割结果 4. 批量测试并计算指标 使用场景： - 测试Fine-tune后的模型效果 - 对比预训练模型和Fine-tune模型 - 实际应用场景的推理 """ import os import sys import torch import numpy as np from PIL import Image import matplotlib.pyplot as plt from pathlib import Path from typing import Optional, Union import torch.nn.functional as F from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor import albumentations as A from albumentations.pytorch import ToTensorV2 # 设置matplotlib中文字体 plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'Arial Unicode MS'] # 用来正常显示中文标签 plt.rcParams['axes.unicode_minus'] = False # 用来正常显示负号 # 添加父目录到路径 sys.path.append(str(Path(__file__).parent.parent / "training")) from config import TrainConfig class SegFormerInference: """SegFormer推理类""" def __init__( self, model_path: Optional[str] = None, pretrained_model: str = "nvidia/segformer-b0-finetuned-ade-512-512", num_classes: int = 2, image_size: int = 256, # ⚠️ 重要：必须与训练时一致！ device: str = "auto" ): """ 初始化推理类 Args: model_path: Fine-tune后的模型路径（.pt文件） 如果为None，则使用预训练模型 pretrained_model: 预训练模型名称（用于加载processor） num_classes: 类别数 image_size: 输入图像尺寸（必须与训练时一致！） device: 设备 ('cpu', 'cuda', 'auto') """ self.num_classes = num_classes self.image_size = image_size # 设置设备 if device == "auto": self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") else: self.device = torch.device(device) print(f"使用设备: {self.device}") # ⚠️ 重要：使用与训练时完全一致的预处理 # 不再使用 SegformerImageProcessor，而是手动构建预处理pipeline print(f"构建预处理Pipeline（与训练时一致）") print(f" 图像尺寸: {image_size}×{image_size}") print(f" 归一化: ImageNet标准（mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]）") self.transform = A.Compose([ A.Resize(image_size, image_size), # ⚠️ 关键：必须resize到训练时的尺寸 A.Normalize( mean=[0.485, 0.456, 0.406], # ImageNet标准均值 std=[0.229, 0.224, 0.225], # ImageNet标准标准差 ), ToTensorV2(), ]) # 加载模型 if model_path: print(f"加载Fine-tune模型: {model_path}") self.model = self._load_finetuned_model(model_path) else: print(f"加载预训练模型: {pretrained_model}") self.model = SegformerForSemanticSegmentation.from_pretrained( pretrained_model, num_labels=num_classes ) self.model.to(self.device) self.model.eval() print("✓ 模型加载完成") def _load_finetuned_model(self, checkpoint_path: str): """ 加载Fine-tune后的模型 Args: checkpoint_path: checkpoint文件路径 Returns: model: 加载权重后的模型 """ from transformers import SegformerConfig # 加载checkpoint checkpoint = torch.load(checkpoint_path, map_location=self.device, weights_only=False) # 从checkpoint中获取配置 if 'config' in checkpoint: config_dict = checkpoint['config'] pretrained_model = config_dict.get('pretrained_model', 'nvidia/segformer-b0-finetuned-ade-512-512') else: pretrained_model = 'nvidia/segformer-b0-finetuned-ade-512-512' # 创建模型配置（不加载预训练权重） config = SegformerConfig.from_pretrained(pretrained_model) config.num_labels = self.num_classes # 用配置创建空模型 model = SegformerForSemanticSegmentation(config) # 加载你训练好的权重 model.load_state_dict(checkpoint['model_state_dict']) # 打印训练信息 print(f" ✓ 模型权重加载成功") if 'epoch' in checkpoint: print(f" 训练轮数: {checkpoint['epoch']}") if 'metrics' in checkpoint: metrics = checkpoint['metrics'] if 'miou' in metrics: print(f" mIoU: {metrics['miou']:.4f}") if 'pixel_accuracy' in metrics: print(f" 像素准确率: {metrics['pixel_accuracy']:.4f}") return model @torch.no_grad() def predict(self, image: Union[str, Image.Image, np.ndarray]) -> np.ndarray: """ 对单张图像进行预测 Args: image: 输入图像 - str: 图像文件路径 - PIL.Image: PIL图像对象 - np.ndarray: numpy数组 (H, W, 3) Returns: pred_mask: 预测的分割mask (H, W)，像素值为类别ID """ # 1. 读取图像并保存原始尺寸 if isinstance(image, str): image = Image.open(image).convert('RGB') original_size = image.size # (W, H) image = np.array(image) elif isinstance(image, np.ndarray): original_size = (image.shape[1], image.shape[0]) # (W, H) else: # PIL.Image original_size = image.size # (W, H) image = np.array(image) # 2. 应用与训练时一致的预处理 # ⚠️ 关键：使用ImageNet标准归一化，与训练时完全一致 transformed = self.transform(image=image) image_tensor = transformed['image'] # (3, H, W) # 3. 添加batch维度 image_tensor = image_tensor.unsqueeze(0).to(self.device) # (1, 3, H, W) # 4. 模型推理 outputs = self.model(pixel_values=image_tensor) logits = outputs.logits # (1, num_classes, H/4, W/4) # 5. 上采样到原始尺寸 logits = F.interpolate( logits, size=(original_size[1], original_size[0]), # (H, W) mode='bilinear', align_corners=False ) # 6. 获取预测类别 pred_mask = logits.argmax(dim=1)[0].cpu().numpy() return pred_mask def visualize_prediction( self, image_path: str, pred_mask: Optional[np.ndarray] = None, save_path: Optional[str] = None, class_names: list = None ): """ 可视化预测结果 Args: image_path: 原始图像路径 pred_mask: 预测mask（如果为None，则重新预测） save_path: 保存路径 class_names: 类别名称列表 """ # 读取原始图像 image = Image.open(image_path).convert('RGB') image_np = np.array(image) # 预测 if pred_mask is None: pred_mask = self.predict(image) # 默认类别名称 if class_names is None: class_names = ['背景', '食物'] # 创建可视化 fig, axes = plt.subplots(1, 3, figsize=(18, 6)) # 1. 原始图像 axes[0].imshow(image_np) axes[0].set_title('原始图像', fontsize=14) axes[0].axis('off') # 2. 预测mask（彩色） axes[1].imshow(pred_mask, cmap='tab20', vmin=0, vmax=self.num_classes-1) axes[1].set_title('预测Mask', fontsize=14) axes[1].axis('off') # 3. 叠加显示 # 创建彩色mask colored_mask = np.zeros_like(image_np) colors = [ [0, 0, 0], # 背景：黑色 [255, 0, 0], # 食物：红色 ] for class_id in range(self.num_classes): colored_mask[pred_mask == class_id] = colors[class_id] # 叠加 alpha = 0.5 overlay = (image_np * (1 - alpha) + colored_mask * alpha).astype(np.uint8) axes[2].imshow(overlay) axes[2].set_title('叠加显示（红色=食物）', fontsize=14) axes[2].axis('off') # 添加统计信息 total_pixels = pred_mask.size food_pixels = np.sum(pred_mask == 1) food_ratio = food_pixels / total_pixels * 100 plt.suptitle( f'分割结果 | 食物区域: {food_ratio:.1f}%', fontsize=16, y=0.98 ) plt.tight_layout() if save_path: plt.savefig(save_path, dpi=150, bbox_inches='tight') print(f"✓ 结果已保存: {save_path}") plt.show() def extract_food_region( self, image_path: str, pred_mask: Optional[np.ndarray] = None, save_path: Optional[str] = None ) -> np.ndarray: """ 提取食物区域（去除背景） Args: image_path: 原始图像路径 pred_mask: 预测mask（如果为None，则重新预测） save_path: 保存路径 Returns: food_region: 提取的食物区域图像（背景为白色） """ # 读取图像 image = Image.open(image_path).convert('RGB') image_np = np.array(image) # 预测 if pred_mask is None: pred_mask = self.predict(image) # 创建二值mask food_mask = (pred_mask == 1) # 提取食物区域（白色背景） food_region = image_np.copy() food_region[~food_mask] = 255 # 保存 if save_path: Image.fromarray(food_region).save(save_path) print(f"✓ 食物区域已保存: {save_path}") return food_region def batch_test( self, image_dir: str, output_dir: str, max_images: int = 10 ): """ 批量测试图像 Args: image_dir: 图像目录 output_dir: 输出目录 max_images: 最大测试图像数 """ image_dir = Path(image_dir) output_dir = Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) # 获取所有图像文件 image_files = [] for ext in ['*.jpg', '*.jpeg', '*.png']: image_files.extend(list(image_dir.glob(ext))) image_files = sorted(image_files)[:max_images] print(f"\n开始批量测试: {len(image_files)} 张图像") print("="*60) for idx, image_path in enumerate(image_files, 1): print(f"[{idx}/{len(image_files)}] {image_path.name}") # 预测 pred_mask = self.predict(str(image_path)) # 可视化 save_path = output_dir / f"{image_path.stem}_result.png" self.visualize_prediction( str(image_path), pred_mask=pred_mask, save_path=str(save_path) ) print("="*60) print(f"✓ 批量测试完成! 结果保存在: {output_dir}") def compare_models( image_path: str, finetuned_model_path: str, pretrained_model: str = "nvidia/segformer-b0-finetuned-ade-512-512" ): """ 对比预训练模型和Fine-tune模型的效果 Args: image_path: 测试图像路径 finetuned_model_path: Fine-tune模型路径 pretrained_model: 预训练模型名称 """ print("\n" + "="*60) print("模型对比测试") print("="*60) # 加载预训练模型 print("\n加载预训练模型...") pretrained_inference = SegFormerInference( model_path=None, pretrained_model=pretrained_model, num_classes=150, # ADE20K的类别数 image_size=256 # 与Fine-tune模型保持一致 ) # 加载Fine-tune模型 print("\n加载Fine-tune模型...") finetuned_inference = SegFormerInference( model_path=finetuned_model_path, pretrained_model=pretrained_model, num_classes=2, image_size=256 # ⚠️ 必须与训练时一致 ) # 读取图像 image = Image.open(image_path).convert('RGB') image_np = np.array(image) # 预测 print("\n执行预测...") pretrained_mask = pretrained_inference.predict(image) finetuned_mask = finetuned_inference.predict(image) # 可视化对比 fig, axes = plt.subplots(2, 3, figsize=(18, 12)) # 预训练模型结果 axes[0, 0].imshow(image_np) axes[0, 0].set_title('原始图像 (预训练模型)', fontsize=12) axes[0, 0].axis('off') axes[0, 1].imshow(pretrained_mask, cmap='tab20b') axes[0, 1].set_title('预训练模型预测\n(ADE20K, 150类)', fontsize=12) axes[0, 1].axis('off') axes[0, 2].imshow(image_np) axes[0, 2].imshow(pretrained_mask, alpha=0.5, cmap='tab20b') axes[0, 2].set_title('预训练模型叠加', fontsize=12) axes[0, 2].axis('off') # Fine-tune模型结果 axes[1, 0].imshow(image_np) axes[1, 0].set_title('原始图像 (Fine-tune模型)', fontsize=12) axes[1, 0].axis('off') axes[1, 1].imshow(finetuned_mask, cmap='gray', vmin=0, vmax=1) axes[1, 1].set_title('Fine-tune模型预测\n(食物分割, 2类)', fontsize=12) axes[1, 1].axis('off') # Fine-tune叠加 overlay = image_np.copy().astype(float) red_mask = np.zeros_like(overlay) red_mask[finetuned_mask == 1] = [255, 0, 0] overlay = (overlay * 0.6 + red_mask * 0.4).astype(np.uint8) axes[1, 2].imshow(overlay) axes[1, 2].set_title('Fine-tune模型叠加\n(红色=食物)', fontsize=12) axes[1, 2].axis('off') plt.suptitle('模型对比：预训练 vs Fine-tune', fontsize=16, y=0.98) plt.tight_layout() save_path = "model_comparison.png" plt.savefig(save_path, dpi=150, bbox_inches='tight') print(f"\n✓ 对比结果已保存: {save_path}") plt.show() def main(): """ 主函数：演示推理用法 """ print("="*60) print("SegFormer模型推理测试") print("="*60) # ==================== 配置区 ==================== # TODO: 根据实际情况修改以下配置 # Fine-tune模型路径（训练完成后的checkpoint） FINETUNED_MODEL_PATH = "d:/MyProjects/PythonProjects/FoodClassifier/SegFormer/outputs/best_model.pt" # 测试图像路径 TEST_IMAGE_PATH = "d:/MyProjects/PythonProjects/FoodClassifier/SegFormer/test_food4.jpg" # 批量测试目录 TEST_IMAGE_DIR = "d:/MyProjects/PythonProjects/FoodClassifier/SegFormer/data/segformer_format/images/val" OUTPUT_DIR = "d:/MyProjects/PythonProjects/FoodClassifier/SegFormer/inference_results" # =============================================== # 检查模型文件是否存在 if not os.path.exists(FINETUNED_MODEL_PATH): print(f"\n⚠️ Fine-tune模型不存在: {FINETUNED_MODEL_PATH}") print(f"\n请先运行训练脚本:") print(f" python training/train_minimal.py") print(f"\n现在将使用预训练模型进行测试...") FINETUNED_MODEL_PATH = None # 创建推理实例 inference = SegFormerInference( model_path=FINETUNED_MODEL_PATH, num_classes=2, image_size=256 # ⚠️ 必须与训练时一致（见config.py第272行） ) # 单张图像测试 if os.path.exists(TEST_IMAGE_PATH): print(f"\n测试图像: {TEST_IMAGE_PATH}") inference.visualize_prediction( TEST_IMAGE_PATH, save_path="test_result.png" ) # 提取食物区域 inference.extract_food_region( TEST_IMAGE_PATH, save_path="food_region.png" ) else: print(f"\n⚠️ 测试图像不存在: {TEST_IMAGE_PATH}") # 批量测试 if os.path.exists(TEST_IMAGE_DIR): print(f"\n批量测试目录: {TEST_IMAGE_DIR}") inference.batch_test( image_dir=TEST_IMAGE_DIR, output_dir=OUTPUT_DIR, max_images=5 ) # 模型对比（如果有Fine-tune模型） if FINETUNED_MODEL_PATH and os.path.exists(FINETUNED_MODEL_PATH) and os.path.exists(TEST_IMAGE_PATH): print(f"\n执行模型对比...") compare_models( image_path=TEST_IMAGE_PATH, finetuned_model_path=FINETUNED_MODEL_PATH ) print("\n" + "="*60) print("✓ 测试完成!") print("="*60) if __name__ == "__main__": main()