From f18063e50fe677e254296f92723f902fbb6257fb Mon Sep 17 00:00:00 2001 From: zhangpu <1250681871@qq.com> Date: Fri, 12 Dec 2025 14:10:21 +0800 Subject: [PATCH] =?UTF-8?q?=E7=A1=AE=E4=BF=9D=E9=A2=84=E5=A4=84=E7=90=86?= =?UTF-8?q?=E6=96=B9=E5=BC=8F=E7=9B=B8=E5=90=8C=E3=80=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- SegFormer/inference/test_model.py | 73 +++++++++++------- .../scripts/1_convert_coco_to_segformer.py | 75 +++++++++++++++++-- 2 files changed, 115 insertions(+), 33 deletions(-) diff --git a/SegFormer/inference/test_model.py b/SegFormer/inference/test_model.py index 8e5a4ca..5770d8c 100644 --- a/SegFormer/inference/test_model.py +++ b/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 diff --git a/SegFormer/scripts/1_convert_coco_to_segformer.py b/SegFormer/scripts/1_convert_coco_to_segformer.py index 033987f..0bc3d57 100644 --- a/SegFormer/scripts/1_convert_coco_to_segformer.py +++ b/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,11 +170,24 @@ class COCOToSegFormerConverter: # 将所有标注合并到同一个mask for ann in annotations: - if 'segmentation' in ann and isinstance(ann['segmentation'], list): - # 转换polygon为mask - obj_mask = self._polygon_to_mask(ann['segmentation'], image_size) - # 合并到总mask（取并集） - merged_mask = np.maximum(merged_mask, obj_mask) + if 'segmentation' not in ann: + continue + + 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) return merged_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"