From 92df030d95535551d6feb1e8cce99f3ce6f12d4a Mon Sep 17 00:00:00 2001 From: zhangpu <1250681871@qq.com> Date: Wed, 10 Dec 2025 14:21:10 +0800 Subject: [PATCH] =?UTF-8?q?=E5=A4=84=E7=90=86=E4=B8=AD=E6=96=87=E6=98=BE?= =?UTF-8?q?=E7=A4=BA=E5=BC=82=E5=B8=B8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- SegFormer/inference/test_model.py | 495 ++++++++++++++++++++++++++++++ 1 file changed, 495 insertions(+) create mode 100644 SegFormer/inference/test_model.py diff --git a/SegFormer/inference/test_model.py b/SegFormer/inference/test_model.py new file mode 100644 index 0000000..8e5a4ca --- /dev/null +++ b/SegFormer/inference/test_model.py @@ -0,0 +1,495 @@ +""" +模型推理测试脚本 + +功能说明： +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 + +# 设置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, + device: str = "auto" + ): + """ + 初始化推理类 + + Args: + model_path: Fine-tune后的模型路径（.pt文件） + 如果为None，则使用预训练模型 + pretrained_model: 预训练模型名称（用于加载processor） + num_classes: 类别数 + device: 设备 ('cpu', 'cuda', 'auto') + """ + self.num_classes = num_classes + + # 设置设备 + 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}") + + # 加载processor + print(f"加载Processor: {pretrained_model}") + self.processor = SegformerImageProcessor.from_pretrained(pretrained_model) + + # 加载模型 + 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: 加载权重后的模型 + """ + # 加载checkpoint + checkpoint = torch.load(checkpoint_path, map_location=self.device ,weights_only=False) + + # 从checkpoint中获取配置 + if 'config' in checkpoint: + config = checkpoint['config'] + pretrained_model = config.get('pretrained_model', 'nvidia/segformer-b0-finetuned-ade-512-512') + else: + pretrained_model = 'nvidia/segformer-b0-finetuned-ade-512-512' + + # 创建模型 + model = SegformerForSemanticSegmentation.from_pretrained( + pretrained_model, + num_labels=self.num_classes, + ignore_mismatched_sizes=True + ) + + # 加载权重 + model.load_state_dict(checkpoint['model_state_dict']) + + # 打印训练信息 + if 'epoch' in checkpoint: + print(f" Epoch: {checkpoint['epoch']}") + if 'metrics' in checkpoint: + metrics = checkpoint['metrics'] + print(f" mIoU: {metrics.get('miou', 'N/A'):.4f}") + print(f" Pixel Acc: {metrics.get('pixel_accuracy', 'N/A'):.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 + elif isinstance(image, np.ndarray): + original_size = (image.shape[1], image.shape[0]) # (W, H) + image = Image.fromarray(image) + else: # PIL.Image + original_size = image.size + + # 2. 使用processor预处理 + inputs = self.processor(images=image, return_tensors="pt") + inputs = {k: v.to(self.device) for k, v in inputs.items()} + + # 3. 模型推理 + outputs = self.model(**inputs) + logits = outputs.logits # (1, num_classes, H/4, W/4) + + # 4. 上采样到原始尺寸 + logits = F.interpolate( + logits, + size=(original_size[1], original_size[0]), # (H, W) + mode='bilinear', + align_corners=False + ) + + # 5. 获取预测类别 + 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的类别数 + ) + + # 加载Fine-tune模型 + print("\n加载Fine-tune模型...") + finetuned_inference = SegFormerInference( + model_path=finetuned_model_path, + pretrained_model=pretrained_model, + num_classes=2 + ) + + # 读取图像 + 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 + ) + + # 单张图像测试 + 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()