Compare commits

...
10 Commits
12 changed files with 3068 additions and 9 deletions
+221
View File
@@ -0,0 +1,221 @@
"""
SegFormer-B0 最小化Demo
用于快速上手图像语义分割
"""
import os
import torch
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
class SegFormerDemo:
def __init__(self, model_name="nvidia/segformer-b0-finetuned-ade-512-512"):
"""
初始化SegFormer-B0模型
Args:
model_name: 预训练模型名称
- "nvidia/segformer-b0-finetuned-ade-512-512" (ADE20K数据集,150类)
- "nvidia/segformer-b0-finetuned-cityscapes-1024-1024" (城市街景)
"""
print(f"加载SegFormer模型: {model_name}")
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {self.device}")
# 加载处理器和模型
self.processor = SegformerImageProcessor.from_pretrained(model_name)
self.model = SegformerForSemanticSegmentation.from_pretrained(model_name)
self.model.to(self.device)
self.model.eval()
# ADE20K数据集的类别数
self.num_classes = self.model.config.num_labels
print(f"模型加载完成! 支持 {self.num_classes} 个类别")
def segment_image(self, image_path):
"""
对图像进行语义分割
Args:
image_path: 图像路径
Returns:
pred_seg: 分割结果 (H, W) 每个像素的类别ID
original_image: 原始图像
"""
# 读取图像
image = Image.open(image_path).convert("RGB")
# 预处理
inputs = self.processor(images=image, return_tensors="pt")
inputs = {k: v.to(self.device) for k, v in inputs.items()}
# 推理
with torch.no_grad():
outputs = self.model(**inputs)
logits = outputs.logits
# 上采样到原始尺寸
upsampled_logits = torch.nn.functional.interpolate(
logits,
size=image.size[::-1], # (height, width)
mode="bilinear",
align_corners=False
)
# 获取每个像素的类别
pred_seg = upsampled_logits.argmax(dim=1)[0].cpu().numpy()
return pred_seg, np.array(image)
def visualize_segmentation(self, image, pred_seg, save_path=None):
"""
可视化分割结果
Args:
image: 原始图像 (H, W, 3)
pred_seg: 分割结果 (H, W)
save_path: 保存路径
"""
fig, axes = plt.subplots(1, 3, figsize=(15, 5))
# 原始图像
axes[0].imshow(image)
axes[0].set_title("Original Image")
axes[0].axis("off")
# 分割mask(彩色)
axes[1].imshow(pred_seg, cmap='tab20b')
axes[1].set_title("Segmentation Mask")
axes[1].axis("off")
# 叠加显示
axes[2].imshow(image)
axes[2].imshow(pred_seg, alpha=0.5, cmap='tab20b')
axes[2].set_title("Overlay")
axes[2].axis("off")
plt.tight_layout()
if save_path:
plt.savefig(save_path, dpi=150, bbox_inches='tight')
print(f"结果已保存到: {save_path}")
plt.show()
def extract_object_mask(self, pred_seg, target_classes):
"""
提取特定类别的mask
Args:
pred_seg: 分割结果 (H, W)
target_classes: 目标类别ID列表
Returns:
mask: 二值mask (H, W)
"""
mask = np.isin(pred_seg, target_classes)
return mask.astype(np.uint8) * 255
def extract_food_region(self, image, pred_seg):
"""
提取食物区域(基于ADE20K数据集)
ADE20K中食物相关类别(部分):
- 53: food, 54: table, 55: chair, 56: car, 57: person
- 116: plate, 117: banana, 118: apple, 119: sandwich, 120: orange
Args:
image: 原始图像 (H, W, 3)
pred_seg: 分割结果 (H, W)
Returns:
food_image: 提取的食物区域图像
food_mask: 食物区域mask
"""
# 食物相关类别(需要根据实际情况调整)
food_classes = [53, 116, 117, 118, 119, 120, 121, 122, 123]
# 提取mask
food_mask = self.extract_object_mask(pred_seg, food_classes)
# 应用mask
food_image = image.copy()
food_image[food_mask == 0] = 255 # 白色背景
return food_image, food_mask
def main():
"""
主函数:演示SegFormer-B0的基本用法
"""
print("=" * 50)
print("SegFormer-B0 快速上手Demo")
print("=" * 50)
# 1. 创建Demo实例
demo = SegFormerDemo()
# 2. 测试图像路径(请替换为您的图像路径)
test_image = "test_food3.jpg"
# 如果没有测试图像,创建一个示例
if not os.path.exists(test_image):
print(f"\n警告: 未找到测试图像 '{test_image}'")
print("请将您的食物图片重命名为 'test_food.jpg' 或修改代码中的路径")
print("\n提示: 您可以从以下位置获取测试图像:")
print(" - dataset目录中的训练图像")
print(" - 任何食物图片")
return
# 3. 进行分割
print(f"\n正在处理图像: {test_image}")
pred_seg, original_image = demo.segment_image(test_image)
# 4. 显示统计信息
unique_classes = np.unique(pred_seg)
print(f"\n检测到的类别数: {len(unique_classes)}")
print(f"类别ID: {unique_classes[:20]}...") # 显示前20个
# 5. 可视化结果
print("\n可视化分割结果...")
demo.visualize_segmentation(
original_image,
pred_seg,
save_path="segmentation_result.png"
)
# 6. 提取食物区域
print("\n提取食物区域...")
food_image, food_mask = demo.extract_food_region(original_image, pred_seg)
# 可视化食物区域
fig, axes = plt.subplots(1, 3, figsize=(15, 5))
axes[0].imshow(original_image)
axes[0].set_title("Original")
axes[0].axis("off")
axes[1].imshow(food_mask, cmap='gray')
axes[1].set_title("Food Mask")
axes[1].axis("off")
axes[2].imshow(food_image)
axes[2].set_title("Extracted Food Region")
axes[2].axis("off")
plt.tight_layout()
plt.savefig("food_extraction_result.png", dpi=150)
print("食物提取结果已保存到: food_extraction_result.png")
plt.show()
print("\n" + "=" * 50)
print("Demo运行完成!")
print("=" * 50)
if __name__ == "__main__":
main()
+524
View File
@@ -0,0 +1,524 @@
"""
模型推理测试脚本
功能说明:
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
import albumentations as A
from albumentations.pytorch import ToTensorV2
# 设置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,
image_size: int = 256, # ⚠️ 重要:必须与训练时一致!
device: str = "auto"
):
"""
初始化推理类
Args:
model_path: Fine-tune后的模型路径(.pt文件)
如果为None,则使用预训练模型
pretrained_model: 预训练模型名称(用于加载processor)
num_classes: 类别数
image_size: 输入图像尺寸(必须与训练时一致!)
device: 设备 ('cpu', 'cuda', 'auto')
"""
self.num_classes = num_classes
self.image_size = image_size
# 设置设备
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}")
# ⚠️ 重要:使用与训练时完全一致的预处理
# 不再使用 SegformerImageProcessor,而是手动构建预处理pipeline
print(f"构建预处理Pipeline(与训练时一致)")
print(f" 图像尺寸: {image_size}×{image_size}")
print(f" 归一化: ImageNet标准(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]")
self.transform = A.Compose([
A.Resize(image_size, image_size), # ⚠️ 关键:必须resize到训练时的尺寸
A.Normalize(
mean=[0.485, 0.456, 0.406], # ImageNet标准均值
std=[0.229, 0.224, 0.225], # ImageNet标准标准差
),
ToTensorV2(),
])
# 加载模型
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: 加载权重后的模型
"""
from transformers import SegformerConfig
# 加载checkpoint
checkpoint = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
# 从checkpoint中获取配置
if 'config' in checkpoint:
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'
# 创建模型配置(不加载预训练权重)
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" 训练轮数: {checkpoint['epoch']}")
if 'metrics' in checkpoint:
metrics = checkpoint['metrics']
if 'miou' in metrics:
print(f" mIoU: {metrics['miou']:.4f}")
if 'pixel_accuracy' in metrics:
print(f" 像素准确率: {metrics['pixel_accuracy']:.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 # (W, H)
image = np.array(image)
elif isinstance(image, np.ndarray):
original_size = (image.shape[1], image.shape[0]) # (W, H)
else: # PIL.Image
original_size = image.size # (W, H)
image = np.array(image)
# 2. 应用与训练时一致的预处理
# ⚠️ 关键:使用ImageNet标准归一化,与训练时完全一致
transformed = self.transform(image=image)
image_tensor = transformed['image'] # (3, H, W)
# 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)
# 5. 上采样到原始尺寸
logits = F.interpolate(
logits,
size=(original_size[1], original_size[0]), # (H, W)
mode='bilinear',
align_corners=False
)
# 6. 获取预测类别
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的类别数
image_size=256 # 与Fine-tune模型保持一致
)
# 加载Fine-tune模型
print("\n加载Fine-tune模型...")
finetuned_inference = SegFormerInference(
model_path=finetuned_model_path,
pretrained_model=pretrained_model,
num_classes=2,
image_size=256 # ⚠️ 必须与训练时一致
)
# 读取图像
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,
image_size=256 # ⚠️ 必须与训练时一致(见config.py第272行)
)
# 单张图像测试
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()
@@ -0,0 +1,366 @@
"""
COCO格式数据转换为SegFormer训练格式
功能说明:
1. 读取CVAT导出的COCO格式标注文件
2. 将polygon标注转换为像素级maskPNG图像)
3. 生成训练集和验证集的划分
4. 输出符合SegFormer训练要求的目录结构
输出目录结构:
data/segformer_format/
├── images/
│ ├── train/
│ │ ├── img1.jpg
│ │ └── img2.jpg
│ └── val/
│ └── img3.jpg
└── annotations/
├── train/
│ ├── img1.png # 像素值:0=背景, 1=食物区域
│ └── img2.png
└── val/
└── img3.png
"""
import os
import json
import numpy as np
from PIL import Image, ImageDraw
from pathlib import Path
from typing import List, Dict, Tuple
import shutil
class COCOToSegFormerConverter:
"""COCO格式到SegFormer格式的转换器"""
def __init__(self, coco_json_path: str, coco_images_dir: str, output_dir: str):
"""
初始化转换器
Args:
coco_json_path: COCO标注JSON文件路径(例如:instances_default.json
coco_images_dir: COCO图像所在目录
output_dir: 输出目录(将创建segformer_format文件夹)
"""
self.coco_json_path = coco_json_path
self.coco_images_dir = coco_images_dir
self.output_dir = Path(output_dir)
# 加载COCO标注
print(f"加载COCO标注文件: {coco_json_path}")
with open(coco_json_path, 'r', encoding='utf-8') as f:
self.coco_data = json.load(f)
print(f"✓ 图像数量: {len(self.coco_data['images'])}")
print(f"✓ 标注数量: {len(self.coco_data['annotations'])}")
print(f"✓ 类别数量: {len(self.coco_data['categories'])}")
# 创建输出目录结构
self._create_output_dirs()
def _create_output_dirs(self):
"""创建输出目录结构"""
dirs = [
self.output_dir / "images" / "train",
self.output_dir / "images" / "val",
self.output_dir / "annotations" / "train",
self.output_dir / "annotations" / "val",
]
for d in dirs:
d.mkdir(parents=True, exist_ok=True)
print(f"✓ 输出目录创建完成: {self.output_dir}")
def _polygon_to_mask(self, segmentation: List, image_size: Tuple[int, int]) -> np.ndarray:
"""
将COCO的polygon格式转换为像素级mask
Args:
segmentation: COCO的segmentation字段(polygon列表)
image_size: 图像尺寸 (width, height)
Returns:
mask: 二值mask数组 (H, W),1表示目标区域,0表示背景
"""
width, height = image_size
mask = Image.new('L', (width, height), 0) # 黑色背景
# COCO的segmentation可能包含多个polygon(例如:一个物体被遮挡分成多个部分)
for polygon in segmentation:
# polygon格式: [x1, y1, x2, y2, x3, y3, ...]
# 转换为坐标点列表: [(x1,y1), (x2,y2), ...]
points = [(polygon[i], polygon[i+1]) for i in range(0, len(polygon), 2)]
# 在mask上绘制填充的多边形
ImageDraw.Draw(mask).polygon(points, outline=1, fill=1)
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
image_size: 图像尺寸 (width, height)
Returns:
merged_mask: 合并后的mask (H, W)
像素值: 0=背景(未标注区域), 1=食物区域
"""
width, height = image_size
merged_mask = np.zeros((height, width), dtype=np.uint8)
# 找到该图像的所有标注
annotations = [ann for ann in self.coco_data['annotations']
if ann['image_id'] == image_id]
# 将所有标注合并到同一个mask
for ann in annotations:
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
def convert(self, train_ratio: float = 0.8, random_seed: int = 42):
"""
执行转换流程
Args:
train_ratio: 训练集占比(0.8表示80%训练,20%验证)
random_seed: 随机种子,确保每次划分一致
"""
print("\n" + "="*60)
print("开始转换数据集")
print("="*60)
# 创建图像ID到文件名的映射
id_to_image = {img['id']: img for img in self.coco_data['images']}
# 随机划分训练集和验证集
np.random.seed(random_seed)
image_ids = list(id_to_image.keys())
np.random.shuffle(image_ids)
split_idx = int(len(image_ids) * train_ratio)
train_ids = image_ids[:split_idx]
val_ids = image_ids[split_idx:]
print(f"\n数据集划分:")
print(f" 训练集: {len(train_ids)} 张图像")
print(f" 验证集: {len(val_ids)} 张图像")
# 处理训练集
print(f"\n处理训练集...")
self._process_split(train_ids, id_to_image, split='train')
# 处理验证集
print(f"\n处理验证集...")
self._process_split(val_ids, id_to_image, split='val')
# 保存数据集统计信息
self._save_dataset_info(train_ids, val_ids)
print("\n" + "="*60)
print("✓ 数据集转换完成!")
print("="*60)
print(f"\n输出目录: {self.output_dir}")
print("\n下一步: 运行 2_visualize_data.py 查看转换结果")
def _process_split(self, image_ids: List[int], id_to_image: Dict, split: str):
"""
处理训练集或验证集
Args:
image_ids: 图像ID列表
id_to_image: ID到图像信息的映射
split: 'train''val'
"""
for idx, img_id in enumerate(image_ids, 1):
img_info = id_to_image[img_id]
file_name = img_info['file_name']
width = img_info['width']
height = img_info['height']
print(f" [{idx}/{len(image_ids)}] {file_name}")
# 1. 复制图像到目标目录
src_image_path = Path(self.coco_images_dir) / file_name
dst_image_path = self.output_dir / "images" / split / file_name
if not src_image_path.exists():
print(f" ⚠️ 警告: 图像文件不存在 {src_image_path}")
continue
shutil.copy2(src_image_path, dst_image_path)
# 2. 生成mask并保存为PNG
mask = self._merge_annotations(img_id, (width, height))
# 保存mask(像素值即为类别ID:0=背景, 1=食物)
mask_filename = Path(file_name).stem + '.png' # 改为.png扩展名
mask_path = self.output_dir / "annotations" / split / mask_filename
# 使用PIL保存,确保像素值不被压缩
Image.fromarray(mask, mode='L').save(mask_path)
# 统计信息
food_pixels = np.sum(mask == 1)
total_pixels = mask.size
food_ratio = food_pixels / total_pixels * 100
print(f" ✓ 食物区域占比: {food_ratio:.1f}%")
def _save_dataset_info(self, train_ids: List[int], val_ids: List[int]):
"""保存数据集统计信息"""
info = {
"dataset_name": "Food Segmentation Dataset",
"num_classes": 2, # 背景 + 食物
"class_names": ["background", "food"],
"train_size": len(train_ids),
"val_size": len(val_ids),
"total_size": len(train_ids) + len(val_ids),
"image_format": "jpg/png",
"annotation_format": "png (pixel value = class id)",
}
info_path = self.output_dir / "dataset_info.json"
with open(info_path, 'w', encoding='utf-8') as f:
json.dump(info, f, indent=2, ensure_ascii=False)
print(f"\n✓ 数据集信息已保存: {info_path}")
def main():
"""
主函数:配置路径并执行转换
使用前请修改以下路径:
1. COCO_JSON_PATH: CVAT导出的JSON文件路径
2. COCO_IMAGES_DIR: CVAT导出的图像目录
3. OUTPUT_DIR: 输出目录(将创建segformer_format文件夹)
"""
# ==================== 配置区 ====================
# 请根据您的实际路径修改以下三个变量
# CVAT导出的COCO标注文件(通常名为instances_default.json
COCO_JSON_PATH = "d:/MyProjects/PythonProjects/FoodClassifier/SegFormer/data/raw_coco/annotations/instances_default.json"
# CVAT导出的图像目录
COCO_IMAGES_DIR = "d:/MyProjects/PythonProjects/FoodClassifier/SegFormer/data/raw_coco/images"
# 输出目录(将在此目录下创建segformer_format文件夹)
OUTPUT_DIR = "d:/MyProjects/PythonProjects/FoodClassifier/SegFormer/data/segformer_format"
# 训练集/验证集划分比例(0.8表示80%训练,20%验证)
TRAIN_RATIO = 0.8
# 随机种子(保证每次运行划分结果一致)
RANDOM_SEED = 42
# ===============================================
print("COCO数据集转换工具")
print("目标格式: SegFormer训练格式\n")
# 检查输入文件是否存在
if not os.path.exists(COCO_JSON_PATH):
print(f"❌ 错误: COCO标注文件不存在")
print(f" 路径: {COCO_JSON_PATH}")
print(f"\n请检查:")
print(f" 1. 是否已从CVAT导出COCO格式数据")
print(f" 2. 标注文件路径是否正确")
return
if not os.path.exists(COCO_IMAGES_DIR):
print(f"❌ 错误: 图像目录不存在")
print(f" 路径: {COCO_IMAGES_DIR}")
return
# 创建转换器并执行转换
converter = COCOToSegFormerConverter(
coco_json_path=COCO_JSON_PATH,
coco_images_dir=COCO_IMAGES_DIR,
output_dir=OUTPUT_DIR
)
converter.convert(train_ratio=TRAIN_RATIO, random_seed=RANDOM_SEED)
print("\n" + "="*60)
print("转换完成! 接下来的步骤:")
print("="*60)
print("1. 运行 2_visualize_data.py 检查转换结果")
print("2. 运行 3_train_minimal.py 开始训练")
if __name__ == "__main__":
main()
+392
View File
@@ -0,0 +1,392 @@
"""
数据集可视化工具
功能说明:
1. 可视化转换后的训练数据
2. 检查图像和mask是否正确对齐
3. 统计数据集的基本信息
4. 帮助发现标注错误
使用场景:
- 转换完成后,首先运行此脚本检查数据质量
- 训练前验证数据加载是否正确
"""
import os
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from pathlib import Path
import random
# 设置matplotlib中文字体
plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'Arial Unicode MS'] # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False # 用来正常显示负号
class DatasetVisualizer:
"""数据集可视化工具"""
def __init__(self, data_root: str):
"""
初始化可视化工具
Args:
data_root: 数据集根目录(包含images和annotations文件夹)
"""
self.data_root = Path(data_root)
self.train_images_dir = self.data_root / "images" / "train"
self.train_masks_dir = self.data_root / "annotations" / "train"
self.val_images_dir = self.data_root / "images" / "val"
self.val_masks_dir = self.data_root / "annotations" / "val"
# 检查目录是否存在
if not self.data_root.exists():
raise FileNotFoundError(f"数据集目录不存在: {self.data_root}")
print(f"✓ 数据集根目录: {self.data_root}")
def get_dataset_stats(self):
"""获取数据集统计信息"""
print("\n" + "="*60)
print("数据集统计信息")
print("="*60)
# 统计图像数量
train_images = list(self.train_images_dir.glob("*.*"))
val_images = list(self.val_images_dir.glob("*.*"))
train_masks = list(self.train_masks_dir.glob("*.png"))
val_masks = list(self.val_masks_dir.glob("*.png"))
print(f"\n训练集:")
print(f" 图像数量: {len(train_images)}")
print(f" 标注数量: {len(train_masks)}")
print(f"\n验证集:")
print(f" 图像数量: {len(val_images)}")
print(f" 标注数量: {len(val_masks)}")
print(f"\n总计:")
print(f" 图像总数: {len(train_images) + len(val_images)}")
print(f" 标注总数: {len(train_masks) + len(val_masks)}")
# 统计mask中的类别分布
if len(train_masks) > 0:
print(f"\n正在分析mask内容...")
self._analyze_masks(train_masks[:3]) # 分析前3个mask
return {
'train': len(train_images),
'val': len(val_images),
'train_masks': len(train_masks),
'val_masks': len(val_masks)
}
def _analyze_masks(self, mask_paths: list):
"""
分析mask的像素值分布
Args:
mask_paths: mask文件路径列表
"""
for mask_path in mask_paths:
mask = np.array(Image.open(mask_path))
unique_values = np.unique(mask)
print(f"\n 文件: {mask_path.name}")
print(f" 尺寸: {mask.shape}")
print(f" 像素值: {unique_values}")
# 统计每个类别的像素数
for val in unique_values:
count = np.sum(mask == val)
ratio = count / mask.size * 100
class_name = "背景" if val == 0 else "食物"
print(f" {class_name}(类别{val}): {count}像素 ({ratio:.2f}%)")
def visualize_samples(self, split='train', num_samples=4, random_selection=True):
"""
可视化数据样本
Args:
split: 'train''val'
num_samples: 要可视化的样本数量
random_selection: 是否随机选择样本
"""
print(f"\n可视化{split}集样本...")
# 获取图像和mask路径
if split == 'train':
images_dir = self.train_images_dir
masks_dir = self.train_masks_dir
else:
images_dir = self.val_images_dir
masks_dir = self.val_masks_dir
# 获取所有图像文件
image_files = sorted(list(images_dir.glob("*.*")))
if len(image_files) == 0:
print(f" ⚠️ {split}集中没有图像文件")
return
# 选择要可视化的样本
if random_selection and len(image_files) > num_samples:
selected_files = random.sample(image_files, num_samples)
else:
selected_files = image_files[:num_samples]
# 创建画布
fig, axes = plt.subplots(num_samples, 3, figsize=(15, 5*num_samples))
if num_samples == 1:
axes = axes.reshape(1, -1)
for idx, image_path in enumerate(selected_files):
# 读取图像
image = Image.open(image_path).convert('RGB')
image_np = np.array(image)
# 读取对应的mask
mask_filename = image_path.stem + '.png'
mask_path = masks_dir / mask_filename
if not mask_path.exists():
print(f" ⚠️ 警告: 找不到mask文件 {mask_path.name}")
continue
mask = np.array(Image.open(mask_path))
# 显示原始图像
axes[idx, 0].imshow(image_np)
axes[idx, 0].set_title(f"原始图像\n{image_path.name}")
axes[idx, 0].axis('off')
# 显示mask(使用不同颜色)
# 0=黑色(背景), 1=白色(食物)
axes[idx, 1].imshow(mask, cmap='gray', vmin=0, vmax=1)
axes[idx, 1].set_title(f"分割Mask\n背景=黑, 食物=白")
axes[idx, 1].axis('off')
# 显示叠加效果
# 创建彩色mask用于叠加显示
colored_mask = np.zeros_like(image_np)
colored_mask[mask == 1] = [255, 0, 0] # 食物区域显示为红色
# 叠加显示
overlay = image_np.copy()
alpha = 0.4 # 透明度
overlay[mask == 1] = (
image_np[mask == 1] * (1 - alpha) +
colored_mask[mask == 1] * alpha
).astype(np.uint8)
axes[idx, 2].imshow(overlay)
axes[idx, 2].set_title("叠加显示\n红色=食物区域")
axes[idx, 2].axis('off')
# 打印统计信息
food_pixels = np.sum(mask == 1)
total_pixels = mask.size
food_ratio = food_pixels / total_pixels * 100
print(f" [{idx+1}] {image_path.name} - 食物占比: {food_ratio:.1f}%")
plt.tight_layout()
# 保存可视化结果
save_path = self.data_root / f"visualization_{split}.png"
plt.savefig(save_path, dpi=120, bbox_inches='tight')
print(f"\n✓ 可视化结果已保存: {save_path}")
plt.show()
def check_data_integrity(self):
"""
检查数据完整性
检查项目:
1. 每张图像是否有对应的mask
2. 图像和mask的尺寸是否匹配
3. mask的像素值是否在有效范围内
"""
print("\n" + "="*60)
print("数据完整性检查")
print("="*60)
issues = []
for split in ['train', 'val']:
print(f"\n检查{split}集...")
if split == 'train':
images_dir = self.train_images_dir
masks_dir = self.train_masks_dir
else:
images_dir = self.val_images_dir
masks_dir = self.val_masks_dir
image_files = list(images_dir.glob("*.*"))
for image_path in image_files:
# 检查1: mask文件是否存在
mask_filename = image_path.stem + '.png'
mask_path = masks_dir / mask_filename
if not mask_path.exists():
issues.append(f"{split}/{image_path.name}: 缺少mask文件")
continue
# 检查2: 尺寸是否匹配
image = Image.open(image_path)
mask = Image.open(mask_path)
if image.size != mask.size:
issues.append(
f"{split}/{image_path.name}: "
f"尺寸不匹配 (图像:{image.size}, mask:{mask.size})"
)
# 检查3: mask像素值是否有效
mask_np = np.array(mask)
unique_values = np.unique(mask_np)
# 有效值应该是0(背景)和1(食物)
invalid_values = [v for v in unique_values if v not in [0, 1]]
if invalid_values:
issues.append(
f"{split}/{image_path.name}: "
f"mask包含无效像素值 {invalid_values}"
)
# 输出检查结果
if len(issues) == 0:
print("\n✓ 数据完整性检查通过! 未发现问题")
else:
print(f"\n⚠️ 发现 {len(issues)} 个问题:")
for issue in issues:
print(f" - {issue}")
return len(issues) == 0
def show_class_distribution(self):
"""
显示类别分布统计
统计整个数据集中背景和食物的像素占比
"""
print("\n" + "="*60)
print("类别分布统计")
print("="*60)
for split in ['train', 'val']:
print(f"\n{split}集:")
if split == 'train':
masks_dir = self.train_masks_dir
else:
masks_dir = self.val_masks_dir
mask_files = list(masks_dir.glob("*.png"))
if len(mask_files) == 0:
print(f" 没有mask文件")
continue
# 统计所有mask的像素分布
total_background = 0
total_food = 0
for mask_path in mask_files:
mask = np.array(Image.open(mask_path))
total_background += np.sum(mask == 0)
total_food += np.sum(mask == 1)
total_pixels = total_background + total_food
print(f" 总像素数: {total_pixels:,}")
print(f" 背景像素: {total_background:,} ({total_background/total_pixels*100:.2f}%)")
print(f" 食物像素: {total_food:,} ({total_food/total_pixels*100:.2f}%)")
print(f" 类别平衡度: {min(total_background, total_food) / max(total_background, total_food):.3f}")
# 绘制饼图
fig, ax = plt.subplots(figsize=(8, 6))
ax.pie(
[total_background, total_food],
labels=['背景', '食物'],
autopct='%1.1f%%',
colors=['#808080', '#FF6B6B'],
startangle=90
)
ax.set_title(f'{split}集 - 类别分布')
save_path = self.data_root / f"class_distribution_{split}.png"
plt.savefig(save_path, dpi=120, bbox_inches='tight')
print(f" ✓ 分布图已保存: {save_path}")
plt.close()
def main():
"""
主函数:运行所有可视化和检查
"""
# ==================== 配置区 ====================
# TODO: 修改为您的数据集路径
DATA_ROOT = "d:/MyProjects/PythonProjects/FoodClassifier/SegFormer/data/segformer_format"
# 可视化参数
NUM_SAMPLES = 3 # 每个集合显示的样本数
RANDOM_SELECTION = False # True=随机选择, False=顺序选择前N个
# ===============================================
print("="*60)
print("数据集可视化工具")
print("="*60)
# 检查数据集路径
if not os.path.exists(DATA_ROOT):
print(f"\n❌ 错误: 数据集目录不存在")
print(f" 路径: {DATA_ROOT}")
print(f"\n请先运行 1_convert_coco_to_segformer.py 转换数据集")
return
# 创建可视化工具
visualizer = DatasetVisualizer(DATA_ROOT)
# 1. 显示数据集统计信息
stats = visualizer.get_dataset_stats()
# 2. 检查数据完整性
is_valid = visualizer.check_data_integrity()
if not is_valid:
print("\n⚠️ 请先修复数据问题再继续训练")
return
# 3. 显示类别分布
visualizer.show_class_distribution()
# 4. 可视化训练集样本
if stats['train'] > 0:
visualizer.visualize_samples(
split='train',
num_samples=min(NUM_SAMPLES, stats['train']),
random_selection=RANDOM_SELECTION
)
# 5. 可视化验证集样本
if stats['val'] > 0:
visualizer.visualize_samples(
split='val',
num_samples=min(NUM_SAMPLES, stats['val']),
random_selection=RANDOM_SELECTION
)
print("\n" + "="*60)
print("✓ 可视化完成!")
print("="*60)
print("\n如果数据没有问题,接下来可以:")
print(" 1. 运行 3_train_minimal.py 开始训练(CPU版本)")
print(" 2. 或运行 4_train_gpu.py 开始训练(GPU版本)")
if __name__ == "__main__":
main()
+317
View File
@@ -0,0 +1,317 @@
"""
训练配置文件
功能说明:
1. 集中管理所有训练超参数
2. 区分CPU和GPU训练配置
3. 便于实验管理和超参数调优
使用方式:
from config import TrainConfig
config = TrainConfig()
# 根据需要修改配置
config.batch_size = 2
"""
from dataclasses import dataclass
from pathlib import Path
import torch
@dataclass
class TrainConfig:
"""
训练配置类
使用dataclass装饰器,自动生成__init__等方法
所有配置都有类型注解和默认值
"""
# ==================== 路径配置 ====================
# 数据集根目录
data_root: str = "d:/MyProjects/PythonProjects/FoodClassifier/SegFormer/data/segformer_format"
# 模型保存目录
output_dir: str = "d:/MyProjects/PythonProjects/FoodClassifier/SegFormer/outputs"
# 预训练模型名称
pretrained_model: str = "nvidia/segformer-b0-finetuned-ade-512-512"
# ==================== 模型配置 ====================
# 类别数(包括背景)
# 0: 背景, 1: 食物
num_classes: int = 2
# 输入图像尺寸
image_size: int = 512
# SegFormer模型版本
# B0: 最轻量 (3.7M参数)
# B1: 轻量 (13.7M)
# B2: 中等 (24.7M)
# B3: 较大 (44.6M)
# B4: 大型 (61.4M)
# B5: 超大 (81.9M)
model_variant: str = "b0"
# ==================== 训练配置 ====================
# 批大小
# CPU训练建议: 1-2
# GPU训练建议: 4-8 (4090可以用8-16)
batch_size: int = 1
# 总训练轮数
num_epochs: int = 50
# 学习率
learning_rate: float = 6e-5
# 权重衰减(L2正则化)
weight_decay: float = 0.01
# 学习率调度器类型
# 'cosine': 余弦退火
# 'linear': 线性衰减
# 'polynomial': 多项式衰减
lr_scheduler_type: str = "cosine"
# Warmup轮数(学习率逐步增加的轮数)
warmup_epochs: int = 5
# 梯度裁剪(防止梯度爆炸)
max_grad_norm: float = 1.0
# ==================== 优化器配置 ====================
# 优化器类型: 'adamw' 或 'sgd'
optimizer_type: str = "adamw"
# AdamW的beta参数
adam_betas: tuple = (0.9, 0.999)
# SGD动量
sgd_momentum: float = 0.9
# ==================== 损失函数配置 ====================
# 是否使用混合损失 (CrossEntropy + Dice)
use_mixed_loss: bool = True
# Dice Loss的权重
dice_loss_weight: float = 0.5
# 类别权重(用于处理类别不平衡)
# None: 自动计算
# List[float]: 手动指定 [背景权重, 食物权重]
class_weights: list = None # 例如: [0.3, 0.7]
# ==================== 数据加载配置 ====================
# 数据加载进程数
# CPU训练: 0 (避免进程间通信开销)
# GPU训练: 4-8
num_workers: int = 0
# 是否使用数据增强
use_augmentation: bool = True
# ==================== 训练策略 ====================
# 是否使用两阶段训练
# Stage 1: 冻结Encoder,只训练Decoder
# Stage 2: Fine-tune整个模型
use_two_stage_training: bool = True
# Stage 1的训练轮数(冻结Encoder
stage1_epochs: int = 10
# ==================== 验证和保存 ====================
# 验证频率(每N个epoch验证一次)
eval_every_n_epochs: int = 5
# 保存checkpoint频率(每N个epoch保存一次)
save_every_n_epochs: int = 10
# 是否只保存最佳模型
save_best_only: bool = True
# 最佳模型的评估指标: 'miou', 'loss', 'pixel_acc'
best_metric: str = "miou"
# ==================== 设备配置 ====================
# 设备: 'cpu', 'cuda', 'auto'
device: str = "auto"
# 是否使用混合精度训练(仅GPU
use_amp: bool = False
# ==================== 日志配置 ====================
# 打印频率(每N个batch打印一次)
print_every_n_batches: int = 5
# 是否保存训练日志
save_logs: bool = True
# 是否启用详细日志(包括每个batch的详细信息)
verbose: bool = True
# ==================== 随机种子 ====================
# 随机种子(确保可复现)
random_seed: int = 42
def __post_init__(self):
"""
初始化后的处理
在所有参数赋值后自动调用,用于:
1. 自动推断设备
2. 创建输出目录
3. 验证配置的合理性
"""
# 1. 自动推断设备
if self.device == "auto":
self.device = "cuda" if torch.cuda.is_available() else "cpu"
# 2. CPU训练时自动调整配置
if self.device == "cpu":
self.use_amp = False # CPU不支持混合精度
self.num_workers = 0 # CPU训练避免多进程开销
if self.batch_size > 2:
print(f"⚠️ CPU训练建议batch_size<=2,当前值 {self.batch_size} 可能很慢")
# 3. 创建输出目录
Path(self.output_dir).mkdir(parents=True, exist_ok=True)
# 4. 验证配置
assert self.num_classes >= 2, "类别数必须>=2"
assert self.batch_size > 0, "batch_size必须>0"
assert self.num_epochs > 0, "num_epochs必须>0"
assert self.learning_rate > 0, "learning_rate必须>0"
if self.use_two_stage_training:
assert self.stage1_epochs < self.num_epochs, \
"stage1_epochs必须小于num_epochs"
def to_dict(self) -> dict:
"""
将配置转换为字典(用于保存)
Returns:
config_dict: 配置字典
"""
return {
k: v for k, v in self.__dict__.items()
if not k.startswith('_')
}
def print_config(self):
"""打印所有配置信息"""
print("\n" + "="*60)
print("训练配置")
print("="*60)
print("\n【路径配置】")
print(f" 数据集: {self.data_root}")
print(f" 输出目录: {self.output_dir}")
print(f" 预训练模型: {self.pretrained_model}")
print("\n【模型配置】")
print(f" 类别数: {self.num_classes}")
print(f" 图像尺寸: {self.image_size}×{self.image_size}")
print(f" 模型版本: SegFormer-{self.model_variant.upper()}")
print("\n【训练配置】")
print(f" 批大小: {self.batch_size}")
print(f" 训练轮数: {self.num_epochs}")
print(f" 学习率: {self.learning_rate}")
print(f" 权重衰减: {self.weight_decay}")
print(f" 学习率调度: {self.lr_scheduler_type}")
print(f" Warmup轮数: {self.warmup_epochs}")
print("\n【训练策略】")
if self.use_two_stage_training:
print(f" 两阶段训练: 是")
print(f" Stage 1 (冻结Encoder): {self.stage1_epochs} epochs")
print(f" Stage 2 (全模型Fine-tune): {self.num_epochs - self.stage1_epochs} epochs")
else:
print(f" 两阶段训练: 否")
print("\n【损失函数】")
if self.use_mixed_loss:
print(f" 混合损失: CrossEntropy + {self.dice_loss_weight}×Dice")
else:
print(f" 损失函数: CrossEntropy")
if self.class_weights:
print(f" 类别权重: {self.class_weights}")
print("\n【设备配置】")
print(f" 设备: {self.device}")
print(f" 混合精度: {'' if self.use_amp else ''}")
print(f" 数据加载进程: {self.num_workers}")
print("\n【验证和保存】")
print(f" 验证频率: 每{self.eval_every_n_epochs}")
print(f" 保存频率: 每{self.save_every_n_epochs}")
print(f" 最佳指标: {self.best_metric}")
print("="*60 + "\n")
def get_cpu_config() -> TrainConfig:
"""
获取CPU训练的推荐配置
适用于:
- 本地开发调试
- 快速验证代码正确性
- 小数据集实验
"""
config = TrainConfig()
# CPU优化配置
config.device = "cpu"
config.batch_size = 1
config.num_workers = 0
config.use_amp = False
config.image_size = 256 # 降低分辨率加快训练
# 快速验证配置
# config.num_epochs = 20
config.num_epochs = 50
config.eval_every_n_epochs = 5
config.save_every_n_epochs = 10
return config
def get_gpu_config() -> TrainConfig:
"""
获取GPU训练的推荐配置
适用于:
- 4090等高性能GPU
- 正式训练
- 追求最佳性能
"""
config = TrainConfig()
# GPU优化配置
config.device = "cuda"
config.batch_size = 8
config.num_workers = 4
config.use_amp = True
config.image_size = 512
# 完整训练配置
config.num_epochs = 100
config.eval_every_n_epochs = 5
config.save_every_n_epochs = 10
return config
if __name__ == "__main__":
"""测试配置"""
print("CPU配置:")
cpu_config = get_cpu_config()
cpu_config.print_config()
print("\n\nGPU配置:")
gpu_config = get_gpu_config()
gpu_config.print_config()
+347
View File
@@ -0,0 +1,347 @@
"""
SegFormer训练数据集类
功能说明:
1. 加载图像和对应的分割mask
2. 数据增强(翻转、旋转、颜色变换等)
3. 预处理和标准化
4. 批量数据加载
"""
import os
import numpy as np
from PIL import Image
import torch
from torch.utils.data import Dataset
from pathlib import Path
from typing import Optional, Tuple, List
import albumentations as A
from albumentations.pytorch import ToTensorV2
class FoodSegmentationDataset(Dataset):
"""
食物分割数据集
数据格式:
- 图像:RGB图像 (.jpg, .png等)
- 标注:PNG格式mask,像素值为类别ID (0=背景, 1=食物)
"""
def __init__(
self,
data_root: str,
split: str = 'train',
image_size: int = 512,
num_classes: int = 2,
augmentation: bool = True
):
"""
初始化数据集
Args:
data_root: 数据集根目录(包含images和annotations文件夹)
split: 'train''val'
image_size: 输入图像尺寸(将resize到此尺寸)
num_classes: 类别数(包括背景)
augmentation: 是否使用数据增强(仅训练集)
"""
self.data_root = Path(data_root)
self.split = split
self.image_size = image_size
self.num_classes = num_classes
self.augmentation = augmentation and (split == 'train')
# 构建图像和mask的路径
self.images_dir = self.data_root / "images" / split
self.masks_dir = self.data_root / "annotations" / split
# 获取所有图像文件
self.image_files = sorted(list(self.images_dir.glob("*.*")))
# 过滤:只保留有对应mask的图像
self.valid_samples = []
for img_path in self.image_files:
mask_path = self.masks_dir / (img_path.stem + '.png')
if mask_path.exists():
self.valid_samples.append((img_path, mask_path))
print(f"{split}集加载完成: {len(self.valid_samples)} 个样本")
# 构建数据增强pipeline
self.transform = self._build_transforms()
def _build_transforms(self):
"""
构建数据增强和预处理pipeline
使用albumentations库进行高效的数据增强
注意:对于分割任务,增强操作需要同时应用到图像和mask
"""
if self.augmentation:
# 训练集:激进的数据增强(因为数据量小)
transform = A.Compose([
# 1. 尺寸调整
A.Resize(self.image_size, self.image_size),
# 2. 几何变换(同时作用于图像和mask)
A.HorizontalFlip(p=0.5), # 50%概率水平翻转
A.VerticalFlip(p=0.3), # 30%概率垂直翻转
A.Rotate(limit=30, p=0.5), # ±30度旋转
A.ShiftScaleRotate(
shift_limit=0.1, # 平移±10%
scale_limit=0.2, # 缩放±20%
rotate_limit=20, # 旋转±20度
p=0.5
),
# 3. 颜色增强(仅作用于图像)
A.RandomBrightnessContrast(
brightness_limit=0.2,
contrast_limit=0.2,
p=0.5
),
A.HueSaturationValue(
hue_shift_limit=20,
sat_shift_limit=30,
val_shift_limit=20,
p=0.5
),
# 4. 模糊和噪声
A.OneOf([
A.GaussianBlur(blur_limit=(3, 5), p=1.0),
A.MedianBlur(blur_limit=5, p=1.0),
], p=0.3),
# 5. 标准化(使用ImageNet均值和标准差)
A.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
),
# 6. 转换为Tensor
ToTensorV2(),
])
else:
# 验证集:仅resize和标准化
transform = A.Compose([
A.Resize(self.image_size, self.image_size),
A.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225],
),
ToTensorV2(),
])
return transform
def __len__(self) -> int:
"""返回数据集大小"""
return len(self.valid_samples)
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
"""
获取一个样本
Args:
idx: 样本索引
Returns:
image: 图像Tensor (3, H, W)
mask: 分割mask Tensor (H, W),像素值为类别ID
"""
img_path, mask_path = self.valid_samples[idx]
# 1. 读取图像和mask
image = np.array(Image.open(img_path).convert('RGB'))
mask = np.array(Image.open(mask_path))
# 2. 应用数据增强
# albumentations会自动将增强同时应用到image和mask
transformed = self.transform(image=image, mask=mask)
image = transformed['image'] # Tensor (3, H, W)
mask = transformed['mask'] # 可能是ndarray或Tensor
# 3. 将mask转换为Tensor(检查类型)
if isinstance(mask, np.ndarray):
mask = torch.from_numpy(mask).long()
else:
mask = mask.long() # 已经是Tensor,直接转换类型
# 4. 检查mask的有效性
# 确保mask的值在[0, num_classes-1]范围内
if mask.max() >= self.num_classes:
print(f"警告: {mask_path.name} 包含无效类别ID {mask.max()}")
mask = torch.clamp(mask, 0, self.num_classes - 1)
return image, mask
def get_sample_info(self, idx: int) -> dict:
"""
获取样本的元信息(用于调试和可视化)
Args:
idx: 样本索引
Returns:
info: 包含文件名、路径等信息的字典
"""
img_path, mask_path = self.valid_samples[idx]
return {
'image_name': img_path.name,
'mask_name': mask_path.name,
'image_path': str(img_path),
'mask_path': str(mask_path),
}
def get_dataloaders(
data_root: str,
batch_size: int = 4,
image_size: int = 512,
num_workers: int = 0,
num_classes: int = 2
) -> Tuple[torch.utils.data.DataLoader, torch.utils.data.DataLoader]:
"""
创建训练集和验证集的DataLoader
Args:
data_root: 数据集根目录
batch_size: 批大小
image_size: 图像尺寸
num_workers: 数据加载进程数(CPU训练时设为0)
num_classes: 类别数
Returns:
train_loader: 训练集DataLoader
val_loader: 验证集DataLoader
"""
# 创建训练集
train_dataset = FoodSegmentationDataset(
data_root=data_root,
split='train',
image_size=image_size,
num_classes=num_classes,
augmentation=True # 训练集使用数据增强
)
# 创建验证集
val_dataset = FoodSegmentationDataset(
data_root=data_root,
split='val',
image_size=image_size,
num_classes=num_classes,
augmentation=False # 验证集不使用数据增强
)
# 创建DataLoader
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=batch_size,
shuffle=True, # 训练集打乱顺序
num_workers=num_workers,
pin_memory=False, # CPU训练时设为False
drop_last=True if len(train_dataset) > batch_size else False
)
val_loader = torch.utils.data.DataLoader(
val_dataset,
batch_size=batch_size,
shuffle=False, # 验证集不打乱
num_workers=num_workers,
pin_memory=False,
drop_last=False
)
print(f"\n✓ DataLoader创建完成")
print(f" 训练集: {len(train_dataset)} 样本, {len(train_loader)} 批次")
print(f" 验证集: {len(val_dataset)} 样本, {len(val_loader)} 批次")
return train_loader, val_loader
def test_dataset():
"""
测试数据集加载是否正常
用于开发调试,验证:
1. 数据能否正确加载
2. 数据增强是否正常工作
3. 数据的shape和类型是否正确
"""
import matplotlib.pyplot as plt
print("="*60)
print("数据集测试")
print("="*60)
# 配置
DATA_ROOT = "d:/MyProjects/PythonProjects/FoodClassifier/SegFormer/data/segformer_format"
# 创建数据集
dataset = FoodSegmentationDataset(
data_root=DATA_ROOT,
split='train',
image_size=512,
augmentation=True
)
# 测试读取第一个样本
image, mask = dataset[0]
info = dataset.get_sample_info(0)
print(f"\n样本信息:")
print(f" 文件名: {info['image_name']}")
print(f" 图像shape: {image.shape}") # 应该是 (3, 512, 512)
print(f" Mask shape: {mask.shape}") # 应该是 (512, 512)
print(f" Mask唯一值: {torch.unique(mask).numpy()}") # 应该是 [0, 1]
print(f" 图像数据范围: [{image.min():.3f}, {image.max():.3f}]")
# 可视化前3个样本(含数据增强效果)
fig, axes = plt.subplots(3, 3, figsize=(12, 12))
for i in range(3):
image, mask = dataset[i]
# 反标准化用于显示
mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
image_denorm = image * std + mean
image_denorm = torch.clamp(image_denorm, 0, 1)
# 转换为numpy用于显示
image_np = image_denorm.permute(1, 2, 0).numpy()
mask_np = mask.numpy()
# 显示图像
axes[i, 0].imshow(image_np)
axes[i, 0].set_title(f"样本{i+1}: 图像")
axes[i, 0].axis('off')
# 显示mask
axes[i, 1].imshow(mask_np, cmap='gray', vmin=0, vmax=1)
axes[i, 1].set_title(f"样本{i+1}: Mask")
axes[i, 1].axis('off')
# 显示叠加
overlay = image_np.copy()
red_mask = np.zeros_like(overlay)
red_mask[mask_np == 1] = [1, 0, 0]
overlay = overlay * 0.6 + red_mask * 0.4
axes[i, 2].imshow(overlay)
axes[i, 2].set_title(f"样本{i+1}: 叠加")
axes[i, 2].axis('off')
plt.tight_layout()
plt.savefig("dataset_test_result.png", dpi=120)
print(f"\n✓ 测试结果已保存: dataset_test_result.png")
plt.show()
print("\n" + "="*60)
print("✓ 数据集测试通过!")
print("="*60)
if __name__ == "__main__":
test_dataset()
+679
View File
@@ -0,0 +1,679 @@
"""
SegFormer Fine-tuning训练脚本(CPU版本)
功能说明:
1. 加载预训练的SegFormer模型
2. 在自定义数据集上Fine-tune
3. 支持两阶段训练(冻结Encoder + 全模型Fine-tune
4. 混合损失函数(CrossEntropy + Dice Loss
5. 完整的训练监控和模型保存
适用场景:
- 本地CPU开发调试
- 验证代码正确性
- 小数据集快速实验
pip install numpy==1.26.4
目前是2.2.6
"""
import os
import sys
import time
import json
import numpy as np
from pathlib import Path
from tqdm import tqdm
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
from transformers import SegformerForSemanticSegmentation
import matplotlib.pyplot as plt
# 导入自定义模块
from config import TrainConfig, get_cpu_config
from dataset import get_dataloaders
class DiceLoss(nn.Module):
"""
Dice Loss for semantic segmentation
Dice系数衡量预测和真实mask的重叠程度:
Dice = 2 * |A ∩ B| / (|A| + |B|)
Dice Loss = 1 - Dice
优点:
- 对类别不平衡不敏感
- 关注重叠区域,适合分割任务
"""
def __init__(self, smooth=1.0):
"""
Args:
smooth: 平滑因子,防止除零
"""
super(DiceLoss, self).__init__()
self.smooth = smooth
def forward(self, logits, targets):
"""
Args:
logits: 模型输出 (B, C, H, W)
targets: 真实标签 (B, H, W)
Returns:
loss: Dice Loss标量
"""
# 1. 将logits转换为概率
probs = F.softmax(logits, dim=1) # (B, C, H, W)
# 2. 将targets转换为one-hot编码
num_classes = logits.shape[1]
targets_one_hot = F.one_hot(targets, num_classes) # (B, H, W, C)
targets_one_hot = targets_one_hot.permute(0, 3, 1, 2).float() # (B, C, H, W)
# 3. 展平空间维度
probs = probs.view(probs.shape[0], probs.shape[1], -1) # (B, C, H*W)
targets_one_hot = targets_one_hot.view(targets_one_hot.shape[0], targets_one_hot.shape[1], -1)
# 4. 计算Dice系数
intersection = (probs * targets_one_hot).sum(dim=2) # (B, C)
union = probs.sum(dim=2) + targets_one_hot.sum(dim=2) # (B, C)
dice = (2.0 * intersection + self.smooth) / (union + self.smooth) # (B, C)
# 5. 对所有类别和batch取平均
dice_loss = 1.0 - dice.mean()
return dice_loss
class MixedLoss(nn.Module):
"""
混合损失:CrossEntropy + Dice Loss
结合两种损失的优点:
- CrossEntropy: 像素级分类,梯度稳定
- Dice Loss: 关注重叠区域,处理类别不平衡
"""
def __init__(self, num_classes, class_weights=None, dice_weight=0.5):
"""
Args:
num_classes: 类别数
class_weights: 类别权重(用于CrossEntropy
dice_weight: Dice Loss的权重
"""
super(MixedLoss, self).__init__()
# CrossEntropy Loss
if class_weights is not None:
class_weights = torch.tensor(class_weights, dtype=torch.float32)
self.ce_loss = nn.CrossEntropyLoss(weight=class_weights)
# Dice Loss
self.dice_loss = DiceLoss()
self.dice_weight = dice_weight
def forward(self, logits, targets):
"""
Args:
logits: 模型输出 (B, C, H, W)
targets: 真实标签 (B, H, W)
Returns:
loss: 混合损失
loss_dict: 各损失分量的字典
"""
ce = self.ce_loss(logits, targets)
dice = self.dice_loss(logits, targets)
total_loss = ce + self.dice_weight * dice
return total_loss, {'ce_loss': ce.item(), 'dice_loss': dice.item()}
class SegmentationMetrics:
"""
分割任务的评估指标
包括:
- mIoU (Mean Intersection over Union)
- Pixel Accuracy
- Per-class IoU
"""
def __init__(self, num_classes):
"""
Args:
num_classes: 类别数
"""
self.num_classes = num_classes
self.reset()
def reset(self):
"""重置所有统计"""
self.confusion_matrix = np.zeros((self.num_classes, self.num_classes))
def update(self, preds, targets):
"""
更新混淆矩阵
Args:
preds: 预测标签 (B, H, W) 或 (N,)
targets: 真实标签 (B, H, W) 或 (N,)
"""
preds = preds.flatten()
targets = targets.flatten()
# 过滤无效值
mask = (targets >= 0) & (targets < self.num_classes)
preds = preds[mask]
targets = targets[mask]
# 更新混淆矩阵
for pred, target in zip(preds, targets):
self.confusion_matrix[target, pred] += 1
def compute_miou(self):
"""
计算mIoU (Mean Intersection over Union)
IoU = TP / (TP + FP + FN)
mIoU = mean(IoU_per_class)
Returns:
miou: 平均IoU
iou_per_class: 每个类别的IoU
"""
# 对角线是TP(预测和真实都是该类别)
tp = np.diag(self.confusion_matrix)
# 每行之和是该类别的真实数量(TP + FN)
# 每列之和是预测为该类别的数量(TP + FP)
fp_fn = self.confusion_matrix.sum(axis=0) + self.confusion_matrix.sum(axis=1) - tp
# IoU = TP / (TP + FP + FN)
iou_per_class = tp / (fp_fn + 1e-10)
# 忽略无效类别(分母为0
valid_classes = fp_fn > 0
miou = iou_per_class[valid_classes].mean()
return miou, iou_per_class
def compute_pixel_accuracy(self):
"""
计算像素准确率
Pixel Accuracy = (TP + TN) / Total
Returns:
pixel_acc: 像素准确率
"""
total_correct = np.diag(self.confusion_matrix).sum()
total_pixels = self.confusion_matrix.sum()
pixel_acc = total_correct / (total_pixels + 1e-10)
return pixel_acc
def get_results(self):
"""
获取所有评估指标
Returns:
results: 包含所有指标的字典
"""
miou, iou_per_class = self.compute_miou()
pixel_acc = self.compute_pixel_accuracy()
return {
'miou': miou,
'pixel_accuracy': pixel_acc,
'iou_per_class': iou_per_class.tolist(),
}
class SegFormerTrainer:
"""SegFormer训练器"""
def __init__(self, config: TrainConfig):
"""
初始化训练器
Args:
config: 训练配置
"""
self.config = config
self.device = torch.device(config.device)
# 设置随机种子
self._set_seed(config.random_seed)
# 创建输出目录
self.output_dir = Path(config.output_dir)
self.output_dir.mkdir(parents=True, exist_ok=True)
# 加载数据
print("\n" + "="*60)
print("加载数据集")
print("="*60)
self.train_loader, self.val_loader = get_dataloaders(
data_root=config.data_root,
batch_size=config.batch_size,
image_size=config.image_size,
num_workers=config.num_workers,
num_classes=config.num_classes
)
# 加载模型
print("\n" + "="*60)
print("加载模型")
print("="*60)
self.model = self._build_model()
# 构建损失函数
self.criterion = self._build_criterion()
# 构建优化器和调度器
self.optimizer = self._build_optimizer()
self.scheduler = self._build_scheduler()
# 训练状态
self.current_epoch = 0
self.best_metric = 0.0
self.train_losses = []
self.val_metrics = []
print("\n" + "="*60)
print("✓ 训练器初始化完成")
print("="*60)
def _set_seed(self, seed):
"""设置随机种子"""
torch.manual_seed(seed)
np.random.seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
def _build_model(self):
"""
构建SegFormer模型
Returns:
model: SegFormer模型
"""
print(f"加载预训练模型: {self.config.pretrained_model}")
# 加载预训练模型
model = SegformerForSemanticSegmentation.from_pretrained(
self.config.pretrained_model,
num_labels=self.config.num_classes,
ignore_mismatched_sizes=True # 允许分类头尺寸不匹配
)
model.to(self.device)
# 统计参数量
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"✓ 模型加载完成")
print(f" 总参数量: {total_params:,}")
print(f" 可训练参数: {trainable_params:,}")
return model
def _build_criterion(self):
"""构建损失函数"""
if self.config.use_mixed_loss:
criterion = MixedLoss(
num_classes=self.config.num_classes,
class_weights=self.config.class_weights,
dice_weight=self.config.dice_loss_weight
)
print(f"✓ 损失函数: Mixed Loss (CE + {self.config.dice_loss_weight}×Dice)")
else:
weights = None
if self.config.class_weights:
weights = torch.tensor(self.config.class_weights, dtype=torch.float32)
criterion = nn.CrossEntropyLoss(weight=weights)
print(f"✓ 损失函数: CrossEntropy")
return criterion.to(self.device)
def _build_optimizer(self):
"""构建优化器"""
if self.config.optimizer_type == 'adamw':
optimizer = AdamW(
self.model.parameters(),
lr=self.config.learning_rate,
weight_decay=self.config.weight_decay,
betas=self.config.adam_betas
)
print(f"✓ 优化器: AdamW (lr={self.config.learning_rate})")
else:
optimizer = torch.optim.SGD(
self.model.parameters(),
lr=self.config.learning_rate,
weight_decay=self.config.weight_decay,
momentum=self.config.sgd_momentum
)
print(f"✓ 优化器: SGD (lr={self.config.learning_rate})")
return optimizer
def _build_scheduler(self):
"""构建学习率调度器"""
if self.config.lr_scheduler_type == 'cosine':
scheduler = CosineAnnealingLR(
self.optimizer,
T_max=self.config.num_epochs,
eta_min=1e-7
)
print(f"✓ 学习率调度: Cosine Annealing")
else:
scheduler = None
print(f"✓ 学习率调度: None")
return scheduler
def _freeze_encoder(self):
"""冻结Encoder(仅训练Decoder"""
for name, param in self.model.named_parameters():
if 'decode_head' not in name:
param.requires_grad = False
trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
print(f" 冻结Encoder,可训练参数: {trainable:,}")
def _unfreeze_all(self):
"""解冻所有参数"""
for param in self.model.parameters():
param.requires_grad = True
trainable = sum(p.numel() for p in self.model.parameters() if p.requires_grad)
print(f" 解冻全模型,可训练参数: {trainable:,}")
def train_epoch(self):
"""
训练一个epoch
Returns:
avg_loss: 平均损失
"""
self.model.train()
epoch_loss = 0.0
epoch_ce_loss = 0.0
epoch_dice_loss = 0.0
pbar = tqdm(self.train_loader, desc=f"Epoch {self.current_epoch}/{self.config.num_epochs}")
for batch_idx, (images, masks) in enumerate(pbar):
# 1. 数据迁移到设备
images = images.to(self.device)
masks = masks.to(self.device)
# 2. 前向传播
outputs = self.model(pixel_values=images)
logits = outputs.logits
# 3. 上采样logits到mask尺寸
logits = F.interpolate(
logits,
size=masks.shape[-2:],
mode='bilinear',
align_corners=False
)
# 4. 计算损失
if isinstance(self.criterion, MixedLoss):
loss, loss_dict = self.criterion(logits, masks)
epoch_ce_loss += loss_dict['ce_loss']
epoch_dice_loss += loss_dict['dice_loss']
else:
loss = self.criterion(logits, masks)
# 5. 反向传播
self.optimizer.zero_grad()
loss.backward()
# 6. 梯度裁剪
if self.config.max_grad_norm > 0:
torch.nn.utils.clip_grad_norm_(
self.model.parameters(),
self.config.max_grad_norm
)
# 7. 更新参数
self.optimizer.step()
# 8. 累计损失
epoch_loss += loss.item()
# 9. 更新进度条
pbar.set_postfix({'loss': f'{loss.item():.4f}'})
# 计算平均损失
avg_loss = epoch_loss / len(self.train_loader)
if isinstance(self.criterion, MixedLoss):
avg_ce = epoch_ce_loss / len(self.train_loader)
avg_dice = epoch_dice_loss / len(self.train_loader)
print(f" Train Loss: {avg_loss:.4f} (CE: {avg_ce:.4f}, Dice: {avg_dice:.4f})")
else:
print(f" Train Loss: {avg_loss:.4f}")
return avg_loss
@torch.no_grad()
def validate(self):
"""
验证模型
Returns:
metrics: 评估指标字典
"""
self.model.eval()
metrics = SegmentationMetrics(self.config.num_classes)
val_loss = 0.0
pbar = tqdm(self.val_loader, desc="Validation")
for images, masks in pbar:
# 1. 数据迁移
images = images.to(self.device)
masks = masks.to(self.device)
# 2. 前向传播
outputs = self.model(pixel_values=images)
logits = outputs.logits
# 3. 上采样
logits = F.interpolate(
logits,
size=masks.shape[-2:],
mode='bilinear',
align_corners=False
)
# 4. 计算损失
if isinstance(self.criterion, MixedLoss):
loss, _ = self.criterion(logits, masks)
else:
loss = self.criterion(logits, masks)
val_loss += loss.item()
# 5. 计算预测
preds = logits.argmax(dim=1)
# 6. 更新指标
metrics.update(preds.cpu().numpy(), masks.cpu().numpy())
# 计算最终指标
results = metrics.get_results()
results['loss'] = val_loss / len(self.val_loader)
print(f"\n Validation Results:")
print(f" Loss: {results['loss']:.4f}")
print(f" mIoU: {results['miou']:.4f}")
print(f" Pixel Acc: {results['pixel_accuracy']:.4f}")
print(f" IoU per class: {[f'{x:.3f}' for x in results['iou_per_class']]}")
return results
def save_checkpoint(self, metrics, is_best=False):
"""
保存checkpoint
Args:
metrics: 评估指标
is_best: 是否是最佳模型
"""
checkpoint = {
'epoch': self.current_epoch,
'model_state_dict': self.model.state_dict(),
'optimizer_state_dict': self.optimizer.state_dict(),
'metrics': metrics,
'config': self.config.to_dict(),
}
# 保存最新checkpoint
checkpoint_path = self.output_dir / f"checkpoint_epoch_{self.current_epoch}.pt"
torch.save(checkpoint, checkpoint_path)
print(f" ✓ Checkpoint已保存: {checkpoint_path}")
# 保存最佳模型
if is_best:
best_path = self.output_dir / "best_model.pt"
torch.save(checkpoint, best_path)
print(f" ✓ 最佳模型已保存: {best_path}")
def train(self):
"""完整训练流程"""
print("\n" + "="*60)
print("开始训练")
print("="*60)
self.config.print_config()
# 两阶段训练
if self.config.use_two_stage_training:
print("\n【Stage 1: 冻结Encoder,训练Decoder】")
self._freeze_encoder()
stage1_end = self.config.stage1_epochs
else:
stage1_end = 0
# 主训练循环
for epoch in range(1, self.config.num_epochs + 1):
self.current_epoch = epoch
# 切换到Stage 2
if self.config.use_two_stage_training and epoch == stage1_end + 1:
print("\n" + "="*60)
print("【Stage 2: Fine-tune全模型】")
print("="*60)
self._unfreeze_all()
# 重建优化器(因为参数变了)
self.optimizer = self._build_optimizer()
self.scheduler = self._build_scheduler()
print(f"\n{'='*60}")
print(f"Epoch {epoch}/{self.config.num_epochs}")
print(f"{'='*60}")
# 训练
train_loss = self.train_epoch()
self.train_losses.append(train_loss)
# 学习率调度
if self.scheduler:
self.scheduler.step()
print(f" Learning Rate: {self.optimizer.param_groups[0]['lr']:.2e}")
# 验证
if epoch % self.config.eval_every_n_epochs == 0:
print(f"\n执行验证...")
val_metrics = self.validate()
self.val_metrics.append((epoch, val_metrics))
# 判断是否是最佳模型
current_metric = val_metrics[self.config.best_metric]
is_best = current_metric > self.best_metric
if is_best:
self.best_metric = current_metric
print(f" 🎉 新的最佳{self.config.best_metric}: {current_metric:.4f}")
# 保存checkpoint
if epoch % self.config.save_every_n_epochs == 0:
if self.config.save_best_only and not is_best:
print(f" 跳过保存(非最佳模型)")
else:
self.save_checkpoint(val_metrics, is_best=is_best)
print("\n" + "="*60)
print("✓ 训练完成!")
print("="*60)
print(f"最佳{self.config.best_metric}: {self.best_metric:.4f}")
# 保存训练曲线
self.plot_training_curves()
def plot_training_curves(self):
"""绘制训练曲线"""
fig, axes = plt.subplots(1, 2, figsize=(15, 5))
# 损失曲线
axes[0].plot(range(1, len(self.train_losses) + 1), self.train_losses)
axes[0].set_xlabel('Epoch')
axes[0].set_ylabel('Training Loss')
axes[0].set_title('Training Loss Curve')
axes[0].grid(True)
# mIoU曲线
if len(self.val_metrics) > 0:
epochs = [x[0] for x in self.val_metrics]
mious = [x[1]['miou'] for x in self.val_metrics]
axes[1].plot(epochs, mious, marker='o')
axes[1].set_xlabel('Epoch')
axes[1].set_ylabel('mIoU')
axes[1].set_title('Validation mIoU Curve')
axes[1].grid(True)
plt.tight_layout()
save_path = self.output_dir / "training_curves.png"
plt.savefig(save_path, dpi=120)
print(f"\n✓ 训练曲线已保存: {save_path}")
plt.close()
def main():
"""主函数"""
print("="*60)
print("SegFormer Fine-tuning训练脚本")
print("="*60)
# 获取配置
config = get_cpu_config()
# 根据需要修改配置
# config.num_epochs = 20
# config.image_size = 256
# config.eval_every_n_epochs = 5
# 创建训练器
trainer = SegFormerTrainer(config)
# 开始训练
trainer.train()
print("\n训练完成! 接下来可以:")
print(" 1. 运行 5_test_model.py 测试模型")
print(" 2. 查看输出目录中的训练曲线和checkpoint")
if __name__ == "__main__":
main()
+1 -1
View File
@@ -67,7 +67,7 @@ class EmbeddingFoodClassifierApp:
# model_path = os.path.join(BASE_DIR, "../model/embedding_20251011_133653/best_embedding_model.pth") # model_path = os.path.join(BASE_DIR, "../model/embedding_20251011_133653/best_embedding_model.pth")
# model_path = os.path.join(BASE_DIR, "../model/ProcessedIngredientRecognition/embedding_20251029_170904/best_embedding_model.pth") # model_path = os.path.join(BASE_DIR, "../model/ProcessedIngredientRecognition/embedding_20251029_170904/best_embedding_model.pth")
# model_path = os.path.join(BASE_DIR, "../model/WholeIngredientRecognition/cosface_20251106_134718/best_cosface_model.pth") # model_path = os.path.join(BASE_DIR, "../model/WholeIngredientRecognition/cosface_20251106_134718/best_cosface_model.pth")
model_path = os.path.join(BASE_DIR, "../model/WholeIngredientRecognition/grid_search_20251127_180944/model_s72.0_m0.35.pth") model_path = os.path.join(BASE_DIR, "../model/WholeIngredientRecognition/grid_search_20251204_140816/model_s64.0_m0.4.pth")
# model_path = os.path.join(BASE_DIR, "../model/DishClassification/cosface_20251111_153649/best_cosface_model.pth") # model_path = os.path.join(BASE_DIR, "../model/DishClassification/cosface_20251111_153649/best_cosface_model.pth")
# FAISS索引目录 # FAISS索引目录
+59 -1
View File
@@ -358,6 +358,64 @@ class FAISSIndexBuilder:
return index return index
def extract_embeddings_only(model_path: str,
train_dir: str,
output_dir: str,
embedding_dim: int = 512,
batch_size: int = 16) -> Tuple[np.ndarray, List[int]]:
"""
仅提取特征向量并保存为JSON(用于可视化)
不构建完整的FAISS索引,节省时间
Args:
model_path: 模型路径
train_dir: 训练数据目录
output_dir: 输出目录
embedding_dim: 特征向量维度
batch_size: 批处理大小
Returns:
(embeddings, labels): 特征向量数组和标签列表
"""
print("=" * 60)
print("开始提取特征向量用于可视化")
print("=" * 60)
builder = FAISSIndexBuilder(model_path, embedding_dim)
# 扫描数据
image_paths, class_names, labels = builder.scan_training_data(train_dir)
builder.image_paths = image_paths
builder.labels = labels
# 提取特征
embeddings = builder.extract_features_batch(image_paths, batch_size)
builder.embeddings = embeddings
# 仅保存 embeddings.json 和 labels.json(可视化需要的)
os.makedirs(output_dir, exist_ok=True)
embeddings_json_path = os.path.join(output_dir, 'embeddings.json')
labels_json_path = os.path.join(output_dir, 'labels.json')
print(f"保存特征向量到: {embeddings_json_path}")
with open(embeddings_json_path, "w", encoding="utf-8") as f:
json.dump(embeddings.tolist(), f, ensure_ascii=False)
print(f"保存标签到: {labels_json_path}")
with open(labels_json_path, "w", encoding="utf-8") as f:
json.dump(labels, f, ensure_ascii=False)
print("=" * 60)
print(f"✓ 特征向量提取完成")
print(f" 输出目录: {output_dir}")
print(f" 样本数: {len(embeddings)}")
print(f" 特征维度: {embedding_dim}")
print("=" * 60)
return embeddings, labels
class FAISSSearcher: class FAISSSearcher:
"""FAISS相似度检索器""" """FAISS相似度检索器"""
@@ -498,7 +556,7 @@ def main():
# MODEL_PATH = "../model/embedding_20251011_133653/best_embedding_model.pth" # MODEL_PATH = "../model/embedding_20251011_133653/best_embedding_model.pth"
# MODEL_PATH = "../model/ProcessedIngredientRecognition/embedding_20251103_172012/best_embedding_model.pth" # MODEL_PATH = "../model/ProcessedIngredientRecognition/embedding_20251103_172012/best_embedding_model.pth"
# MODEL_PATH = "../model/WholeIngredientRecognition/cosface_20251113_160103/best_cosface_model.pth" # MODEL_PATH = "../model/WholeIngredientRecognition/cosface_20251113_160103/best_cosface_model.pth"
MODEL_PATH = "../model/WholeIngredientRecognition/grid_search_20251127_180944/model_s72.0_m0.35.pth" MODEL_PATH = "../model/WholeIngredientRecognition/grid_search_20251204_140816/model_s64.0_m0.4.pth"
# MODEL_PATH = "../model/DishClassification/grid_search_20251121_102723/model_s56.0_m0.4.pth" # MODEL_PATH = "../model/DishClassification/grid_search_20251121_102723/model_s56.0_m0.4.pth"
# MODEL_PATH = "../model/DishClassification/cosface_20251111_153649/best_cosface_model.pth" # MODEL_PATH = "../model/DishClassification/cosface_20251111_153649/best_cosface_model.pth"
# TRAIN_DIR = "../dataset/ProcessedIngredientRecognition/train" # TRAIN_DIR = "../dataset/ProcessedIngredientRecognition/train"
+90 -2
View File
@@ -138,7 +138,17 @@ def reduce_dim(X: np.ndarray, method: str, seed: int, tsne_perplexity: int, umap
raise ValueError(f"未知降维方法: {method}") raise ValueError(f"未知降维方法: {method}")
def plot_2d(Z: np.ndarray, y: np.ndarray, title: str, out_path: Optional[str] = None) -> None: def plot_2d(Z: np.ndarray, y: np.ndarray, title: str, out_path: Optional[str] = None, show: bool = True) -> None:
"""
绘制2D降维散点图
Args:
Z: 降维后的2D坐标
y: 类别标签
title: 图表标题
out_path: 保存路径(可选)
show: 是否显示窗口(默认True,训练脚本中应设为False)
"""
plt.figure(figsize=(8, 7), dpi=120) plt.figure(figsize=(8, 7), dpi=120)
classes = np.unique(y) classes = np.unique(y)
# 构建调色板 # 构建调色板
@@ -157,7 +167,85 @@ def plot_2d(Z: np.ndarray, y: np.ndarray, title: str, out_path: Optional[str] =
plt.tight_layout() plt.tight_layout()
if out_path: if out_path:
plt.savefig(out_path, bbox_inches="tight") plt.savefig(out_path, bbox_inches="tight")
plt.show() if show:
plt.show()
def visualize_embeddings_from_files(embeddings_path: str,
labels_path: str,
output_dir: str,
method: str = "pca",
max_points: Optional[int] = None,
seed: int = 42,
**kwargs) -> str:
"""
从文件加载并可视化embeddings(用于训练脚本调用)
Args:
embeddings_path: embeddings.json路径
labels_path: labels.json路径
output_dir: 输出目录
method: 降维方法 (pca/tsne/umap)
max_points: 抽样上限
seed: 随机种子
**kwargs: 其他降维参数
Returns:
输出的PNG图片路径
"""
print("=" * 60)
print(f"开始可视化 embeddings")
print(f"降维方法: {method}")
print("=" * 60)
# 加载数据
print(f"加载数据: {embeddings_path}")
X = load_embeddings_json(embeddings_path)
y = load_labels_json(labels_path)
# 数据对齐
if X.shape[0] != y.shape[0]:
n = min(X.shape[0], y.shape[0])
print(f"警告: 数据不一致,截断到 {n}")
X, y = X[:n], y[:n]
# 抽样
Xs, ys, _ = subsample(X, y, max_points, seed=seed)
if Xs.shape[0] < X.shape[0]:
print(f"已抽样: {Xs.shape[0]}/{X.shape[0]}")
# 降维参数
tsne_perplexity = kwargs.get('tsne_perplexity', 30)
umap_n_neighbors = kwargs.get('umap_n_neighbors', 15)
umap_min_dist = kwargs.get('umap_min_dist', 0.1)
# 降维
print(f"执行降维: {method}")
Z = reduce_dim(Xs, method, seed, tsne_perplexity, umap_n_neighbors, umap_min_dist)
# 诊断
print("\n==== 诊断信息 ====")
diag_info = diagnostics(Xs, ys, reduced2d=Z, method=method)
print(diag_info)
# 保存诊断信息
diag_path = os.path.join(output_dir, f"embedding_{method}_diagnostics.txt")
with open(diag_path, 'w', encoding='utf-8') as f:
f.write(diag_info)
print(f"✓ 诊断信息已保存: {diag_path}")
# 绘图
out_png = os.path.join(output_dir, f"embedding_{method}_2d.png")
title = f"Embedding {method.upper()} 2D (N={Xs.shape[0]})"
# 训练脚本调用时不显示窗口,只保存图片
plot_2d(Z, ys, title, out_path=out_png, show=False)
plt.close('all') # 关闭所有图形
print(f"✓ 可视化图已保存: {out_png}")
print("=" * 60)
return out_png
def main(): def main():
+2 -2
View File
@@ -19,7 +19,7 @@ def main():
base_model = create_mobile_resnet50_embedding(embedding_dim=512, pretrained=True) base_model = create_mobile_resnet50_embedding(embedding_dim=512, pretrained=True)
# model_path = "../model/DishClassification/grid_search_20251121_102723/model_s56.0_m0.4.pth" # model_path = "../model/DishClassification/grid_search_20251121_102723/model_s56.0_m0.4.pth"
# model_path = "../model/WholeIngredientRecognition/cosface_20251106_134718/best_cosface_model.pth" # model_path = "../model/WholeIngredientRecognition/cosface_20251106_134718/best_cosface_model.pth"
model_path = "../model/WholeIngredientRecognition/grid_search_20251127_180944/model_s72.0_m0.35.pth" model_path = "../model/WholeIngredientRecognition/grid_search_20251204_140816/model_s64.0_m0.4.pth"
# model_path = "../model/ProcessedIngredientRecognition/embedding_20251029_173607/best_embedding_model.pth" # model_path = "../model/ProcessedIngredientRecognition/embedding_20251029_173607/best_embedding_model.pth"
if not os.path.exists(model_path): if not os.path.exists(model_path):
@@ -94,7 +94,7 @@ def main():
# 保存模型,一定不要乱改,防止模型覆盖!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! # 保存模型,一定不要乱改,防止模型覆盖!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
# output_path = "../model/DishClassification/grid_search_20251121_102723/best_embedding_model_mobile.pt" # output_path = "../model/DishClassification/grid_search_20251121_102723/best_embedding_model_mobile.pt"
output_path = "../model/WholeIngredientRecognition/grid_search_20251127_180944/best_embedding_model_mobile.pt" output_path = "../model/WholeIngredientRecognition/grid_search_20251204_140816/best_embedding_model_mobile.pt"
# output_path = "../model/WholeIngredientRecognition/cosface_20251106_134718/best_embedding_model_mobile.pt" # output_path = "../model/WholeIngredientRecognition/cosface_20251106_134718/best_embedding_model_mobile.pt"
# output_path = "../model/ProcessedIngredientRecognition/embedding_20251029_173607/best_embedding_model_mobile.pt" # output_path = "../model/ProcessedIngredientRecognition/embedding_20251029_173607/best_embedding_model_mobile.pt"
traced_model.save(output_path) traced_model.save(output_path)
+70 -3
View File
@@ -46,6 +46,10 @@ logger = logging.getLogger(__name__)
's': [56.0, 60.0, 64.0, 68.0], # scale参数 's': [56.0, 60.0, 64.0, 68.0], # scale参数
'm': [0.32, 0.35, 0.38, 0.40], # margin参数 'm': [0.32, 0.35, 0.38, 0.40], # margin参数
}, },
'whole_ingredient': {
's': [56.0, 60.0, 64.0, 68.0, 72.0],
'm': [0.32, 0.35, 0.38, 0.40, 0.45],
},
""" """
GRID_PARAMS = { GRID_PARAMS = {
'dish': { 'dish': {
@@ -53,8 +57,8 @@ GRID_PARAMS = {
'm': [0.32, 0.35, 0.38, 0.40,0.45,0.50], # margin参数 'm': [0.32, 0.35, 0.38, 0.40,0.45,0.50], # margin参数
}, },
'whole_ingredient': { 'whole_ingredient': {
's': [56.0, 60.0, 64.0, 68.0, 72.0], 's': [64.0, 68.0],
'm': [0.32, 0.35, 0.38, 0.40, 0.45], 'm': [0.40, 0.42, 0.43, 0.45],
}, },
'processed_ingredient': { 'processed_ingredient': {
's': [56.0, 60.0, 64.0, 68.0], 's': [56.0, 60.0, 64.0, 68.0],
@@ -103,6 +107,7 @@ def train_single_config(
val_loader: DataLoader, val_loader: DataLoader,
test_loader: DataLoader, test_loader: DataLoader,
num_classes: int, num_classes: int,
save_dir: str,
max_epochs: int = 100, max_epochs: int = 100,
patience: int = 10, patience: int = 10,
min_epochs: int = 25, min_epochs: int = 25,
@@ -138,6 +143,7 @@ def train_single_config(
best_val_loss = float('inf') best_val_loss = float('inf')
early_stop_counter = 0 early_stop_counter = 0
best_epoch = 0 best_epoch = 0
early_stopped = False # 标记是否触发早停
# 训练历史 # 训练历史
train_losses = [] train_losses = []
@@ -172,6 +178,66 @@ def train_single_config(
# 早停判断 # 早停判断
if epoch >= min_epochs and early_stop_counter >= patience: if epoch >= min_epochs and early_stop_counter >= patience:
logger.info(f'早停触发于Epoch {epoch+1}, 最佳Epoch: {best_epoch+1}') logger.info(f'早停触发于Epoch {epoch+1}, 最佳Epoch: {best_epoch+1}')
early_stopped = True
# ===== 保存最佳模型 =====
best_model_path = os.path.join(save_dir, f'best_model_s{s}_m{m}.pth')
torch.save({
'backbone_state_dict': model.state_dict(),
'head_state_dict': head.state_dict(),
's': s,
'm': m,
'best_val_loss': best_val_loss,
'best_epoch': best_epoch,
}, best_model_path)
logger.info(f'✓ 最佳模型已保存: {best_model_path}')
# ===== 生成特征向量可视化 =====
logger.info('🎨 开始生成特征向量可视化...')
try:
# 动态导入函数
faiss_db_dir = os.path.join(settings.BASE_DIR, 'faiss_vector_db')
if faiss_db_dir not in sys.path:
sys.path.insert(0, faiss_db_dir)
from faiss_vector_db.build_faiss_index import extract_embeddings_only
from faiss_vector_db.visualize_embeddings import visualize_embeddings_from_files
# 为每个配置创建独立的可视化目录
vis_output_dir = os.path.join(save_dir, f'model_s{s}_m{m}_visualization')
# 1. 提取特征向量
logger.info(' 步骤1/2: 提取训练集特征向量...')
extract_embeddings_only(
model_path=best_model_path,
train_dir=cfg.train_dir,
output_dir=vis_output_dir,
embedding_dim=cfg.embedding_dim,
batch_size=16 # 使用较小的batch_size加快速度
)
# 2. 生成可视化(使用PCA方法,快速)
logger.info(' 步骤2/2: 生成可视化图...')
embeddings_json = os.path.join(vis_output_dir, 'embeddings.json')
labels_json = os.path.join(vis_output_dir, 'labels.json')
visualize_embeddings_from_files(
embeddings_path=embeddings_json,
labels_path=labels_json,
output_dir=vis_output_dir,
method='pca', # 使用PCA方法(快速)
max_points=None, # 训练集全量可视化
seed=42
)
logger.info(f'✓ 可视化完成,保存至: {vis_output_dir}')
except Exception as e:
logger.error(f'⚠ 可视化生成失败: {e}')
import traceback
traceback.print_exc()
# ===== 可视化逻辑结束 =====
break break
elapsed_time = time.time() - start_time elapsed_time = time.time() - start_time
@@ -186,7 +252,7 @@ def train_single_config(
's': s, 's': s,
'm': m, 'm': m,
'actual_epochs': actual_epochs, 'actual_epochs': actual_epochs,
'early_stopped': actual_epochs < max_epochs, 'early_stopped': early_stopped,
'training_time_minutes': elapsed_time / 60, 'training_time_minutes': elapsed_time / 60,
'best_val_loss': best_val_loss, 'best_val_loss': best_val_loss,
'final_train_loss': train_losses[-1], 'final_train_loss': train_losses[-1],
@@ -298,6 +364,7 @@ def grid_search_main(
cfg, s, m, cfg, s, m,
train_loader, val_loader, test_loader, train_loader, val_loader, test_loader,
num_classes, num_classes,
save_dir=save_dir,
max_epochs=max_epochs, max_epochs=max_epochs,
patience=patience, patience=patience,
min_epochs=min_epochs min_epochs=min_epochs