确保预处理方式相同。

SegFormer/inference/test_model.py

@@ -23,6 +23,8 @@ 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']  # 用来正常显示中文标签
@@ -63,9 +65,16 @@ class SegFormerInference:
         
         print(f"使用设备: {self.device}")
         
-        # 加载processor
-        print(f"加载Processor: {pretrained_model}")
-        self.processor = SegformerImageProcessor.from_pretrained(pretrained_model)
+        # ⚠️ 重要：使用与训练时完全一致的预处理
+        # 不再使用 SegformerImageProcessor，而是手动构建预处理pipeline
+        print(f"构建预处理Pipeline（与训练时一致）")
+        self.transform = A.Compose([
+            A.Normalize(
+                mean=[0.485, 0.456, 0.406],  # ImageNet标准均值
+                std=[0.229, 0.224, 0.225],   # ImageNet标准标准差
+            ),
+            ToTensorV2(),
+        ])
         
         # 加载模型
         if model_path:
@@ -93,33 +102,38 @@ class SegFormerInference:
         Returns:
             model: 加载权重后的模型
         """
+        from transformers import SegformerConfig
+        
         # 加载checkpoint
-        checkpoint = torch.load(checkpoint_path, map_location=self.device ,weights_only=False)
+        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')
+            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'
         
-        # 创建模型
-        model = SegformerForSemanticSegmentation.from_pretrained(
-            pretrained_model,
-            num_labels=self.num_classes,
-            ignore_mismatched_sizes=True
-        )
+        # 创建模型配置（不加载预训练权重）
+        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"  Epoch: {checkpoint['epoch']}")
+            print(f"  训练轮数: {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}")
+            if 'miou' in metrics:
+                print(f"  mIoU: {metrics['miou']:.4f}")
+            if 'pixel_accuracy' in metrics:
+                print(f"  像素准确率: {metrics['pixel_accuracy']:.4f}")
         
         return model
     
@@ -137,25 +151,30 @@ class SegFormerInference:
         Returns:
             pred_mask: 预测的分割mask (H, W)，像素值为类别ID
         """
-        # 1. 读取和预处理图像
+        # 1. 读取图像并保存原始尺寸
         if isinstance(image, str):
             image = Image.open(image).convert('RGB')
-            original_size = image.size
+            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)
-            image = Image.fromarray(image)
         else:  # PIL.Image
-            original_size = image.size
+            original_size = image.size  # (W, H)
+            image = np.array(image)
         
-        # 2. 使用processor预处理
-        inputs = self.processor(images=image, return_tensors="pt")
-        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+        # 2. 应用与训练时一致的预处理
+        # ⚠️ 关键：使用ImageNet标准归一化，与训练时完全一致
+        transformed = self.transform(image=image)
+        image_tensor = transformed['image']  # (3, H, W)
         
-        # 3. 模型推理
-        outputs = self.model(**inputs)
+        # 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)
         
-        # 4. 上采样到原始尺寸
+        # 5. 上采样到原始尺寸
         logits = F.interpolate(
             logits,
             size=(original_size[1], original_size[0]),  # (H, W)
@@ -163,7 +182,7 @@ class SegFormerInference:
             align_corners=False
         )
         
-        # 5. 获取预测类别
+        # 6. 获取预测类别
         pred_mask = logits.argmax(dim=1)[0].cpu().numpy()
         
         return pred_mask