确保预处理方式相同。

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}")
+        # 不再使用 SegformerImageProcessor，而是手动构建预处理pipeline
-        self.processor = SegformerImageProcessor.from_pretrained(pretrained_model)
+        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中获取配置
         if 'config' in checkpoint:
-            config = checkpoint['config']
+            config_dict = checkpoint['config']
-            pretrained_model = config.get('pretrained_model', 'nvidia/segformer-b0-finetuned-ade-512-512')
+            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(
+        config = SegformerConfig.from_pretrained(pretrained_model)
-            pretrained_model,
+        config.num_labels = self.num_classes
-            num_labels=self.num_classes,
-            ignore_mismatched_sizes=True
-        )
         
-        # 加载权重
+        # 用配置创建空模型
+        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}")
+            if 'miou' in metrics:
-            print(f"  Pixel Acc: {metrics.get('pixel_accuracy', 'N/A'):.4f}")
+                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预处理
+        # 2. 应用与训练时一致的预处理
-        inputs = self.processor(images=image, return_tensors="pt")
+        # ⚠️ 关键：使用ImageNet标准归一化，与训练时完全一致
-        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+        transformed = self.transform(image=image)
+        image_tensor = transformed['image']  # (3, H, W)
         
-        # 3. 模型推理
+        # 3. 添加batch维度
-        outputs = self.model(**inputs)
+        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

SegFormer/scripts/1_convert_coco_to_segformer.py

@@ -97,11 +97,61 @@ class COCOToSegFormerConverter:
         
         return np.array(mask)
     
+    def _rle_to_mask(self, rle: Dict, image_size: Tuple[int, int]) -> np.ndarray:
+        """
+        将COCO的RLE (Run-Length Encoding) 格式转换为像素级mask
+        
+        RLE格式说明：
+        - counts: [v1, v2, v3, ...] 交替表示背景和前景的像素数量
+        - 从背景(0)开始，然后是前景(1)，再是背景(0)，依此类推
+        - 像素顺序：从左到右，从上到下（行优先）
+        
+        Args:
+            rle: COCO的RLE字典，格式为 {"counts": [n1, n2, ...], "size": [height, width]}
+            image_size: 图像尺寸 (width, height)
+        
+        Returns:
+            mask: 二值mask数组 (H, W)，1表示目标区域，0表示背景
+        """
+        width, height = image_size
+        
+        # 获取RLE的counts
+        if 'counts' in rle:
+            counts = rle['counts']
+        else:
+            print(f"    ⚠️  警告: RLE格式不包含counts字段")
+            return np.zeros((height, width), dtype=np.uint8)
+        
+        # 创建一维数组
+        mask = np.zeros(height * width, dtype=np.uint8)
+        
+        current_pos = 0
+        current_value = 0  # 从背景开始
+        
+        for count in counts:
+            # 如果当前是前景(1)，填充mask
+            if current_value == 1:
+                mask[current_pos:current_pos + count] = 1
+            
+            # 移动位置
+            current_pos += count
+            
+            # 切换值（0 -> 1 -> 0 -> 1 ...）
+            current_value = 1 - current_value
+        
+        # 重塑为2D数组 (H, W)
+        mask = mask.reshape((height, width))
+        
+        return mask
+    
     def _merge_annotations(self, image_id: int, image_size: Tuple[int, int]) -> np.ndarray:
         """
         合并一张图像的所有标注为单一mask
         
         由于用户标注时所有食材区域都是同一类别，我们需要将同一图像的多个标注合并
+        支持两种格式：
+        1. Polygon格式: segmentation 是列表的列表
+        2. RLE格式: segmentation 是包含 'counts' 字段的字典
         
         Args:
             image_id: COCO图像ID
@@ -120,9 +170,22 @@ class COCOToSegFormerConverter:
         
         # 将所有标注合并到同一个mask
         for ann in annotations:
-            if 'segmentation' in ann and isinstance(ann['segmentation'], list):
+            if 'segmentation' not in ann:
-                # 转换polygon为mask
+                continue
-                obj_mask = self._polygon_to_mask(ann['segmentation'], image_size)
+            
+            segmentation = ann['segmentation']
+            
+            # 判断是RLE格式还是Polygon格式
+            if isinstance(segmentation, dict) and 'counts' in segmentation:
+                # RLE格式
+                obj_mask = self._rle_to_mask(segmentation, image_size)
+            elif isinstance(segmentation, list) and len(segmentation) > 0:
+                # Polygon格式
+                obj_mask = self._polygon_to_mask(segmentation, image_size)
+            else:
+                print(f"    ⚠️  警告: 未知的segmentation格式: {type(segmentation)}")
+                continue
+            
             # 合并到总mask（取并集）
             merged_mask = np.maximum(merged_mask, obj_mask)
         
@@ -247,7 +310,7 @@ def main():
     """
     
     # ==================== 配置区 ====================
-    # TODO: 请根据您的实际路径修改以下三个变量
+    # 请根据您的实际路径修改以下三个变量
     
     # CVAT导出的COCO标注文件（通常名为instances_default.json）
     COCO_JSON_PATH = "d:/MyProjects/PythonProjects/FoodClassifier/SegFormer/data/raw_coco/annotations/instances_default.json"