确保预处理方式相同。
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user