Mask R-CNN instance segmentation training code
This commit is contained in:
flk
+38
@@ -0,0 +1,38 @@
|
||||
# Python
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
*.so
|
||||
*.egg-info/
|
||||
dist/
|
||||
build/
|
||||
*.egg
|
||||
|
||||
# 虚拟环境
|
||||
env/
|
||||
venv/
|
||||
.venv/
|
||||
|
||||
# PyTorch 模型权重（体积大，不纳入版本控制）
|
||||
model/*.pth
|
||||
checkpoints/*.pth
|
||||
|
||||
# 训练输出
|
||||
checkpoints/loss_curve.png
|
||||
checkpoints/loss_history.json
|
||||
checkpoints/label_map.json
|
||||
predictions/
|
||||
|
||||
# IDE
|
||||
.idea/
|
||||
.vscode/
|
||||
*.swp
|
||||
*.swo
|
||||
|
||||
# OS
|
||||
Thumbs.db
|
||||
Desktop.ini
|
||||
.DS_Store
|
||||
|
||||
# 日志
|
||||
*.log
|
||||
@@ -0,0 +1,161 @@
|
||||
# Mask R-CNN 实例分割训练框架
|
||||
|
||||
基于 PyTorch + torchvision 的 Mask R-CNN 目标分割训练代码，支持 LabelMe 标注格式。
|
||||
|
||||
## 功能特点
|
||||
|
||||
- **LabelMe 标注支持**：自动解析 `polygon`、`linestrip`、`rectangle`、`circle` 四种标注类型
|
||||
- **Linestrip 自动合并**：将相同标签的多条线段智能合并为闭合多边形（适用于标注边界线而非直接画多边形的场景）
|
||||
- **COCO 预训练微调**：加载 COCO 预训练权重，仅替换分类头和 mask 预测头，加速收敛
|
||||
- **数据增强**：随机水平翻转、亮度调整、对比度调整（同时变换 mask 和 boxes）
|
||||
- **训练管理**：检查点保存/恢复、最佳模型追踪、损失曲线可视化
|
||||
- **推理可视化**：支持单图/批量推理，叠加显示 mask 和包围框
|
||||
|
||||
## 环境要求
|
||||
|
||||
- Python 3.11
|
||||
- Windows / Linux / macOS
|
||||
|
||||
## 安装依赖
|
||||
|
||||
```bash
|
||||
# 激活虚拟环境（如已有）
|
||||
# .\env\Scripts\activate
|
||||
|
||||
# 安装 PyTorch (CPU 版本)
|
||||
pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
|
||||
|
||||
# 安装其他依赖
|
||||
pip install opencv-python tqdm
|
||||
|
||||
# 或一次性安装
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
## 项目结构
|
||||
|
||||
```
|
||||
segmentation/
|
||||
├── dataset/ # 数据集目录
|
||||
│ ├── 1.jpg # 图像
|
||||
│ ├── 1.json # LabelMe 标注
|
||||
│ ├── 2.jpg
|
||||
│ ├── 2.json
|
||||
│ └── ...
|
||||
├── utils.py # 工具函数（标注解析、mask 转换、可视化）
|
||||
├── dataset.py # Dataset 类和数据增强
|
||||
├── model.py # Mask R-CNN 模型构建
|
||||
├── train.py # 训练脚本
|
||||
├── predict.py # 推理脚本
|
||||
├── requirements.txt # 依赖列表
|
||||
└── README.md # 本文件
|
||||
```
|
||||
|
||||
## 使用方法
|
||||
|
||||
### 1. 训练模型
|
||||
|
||||
```bash
|
||||
# 默认参数训练（使用 COCO 预训练权重）
|
||||
python train.py
|
||||
|
||||
# 自定义参数
|
||||
python train.py --epochs 100 --batch-size 2 --lr 0.005
|
||||
|
||||
# 冻结骨干网络（小数据集推荐，减少过拟合）
|
||||
python train.py --freeze-backbone --epochs 100
|
||||
|
||||
# 调整图像输入尺寸
|
||||
python train.py --min-size 512 --max-size 800
|
||||
|
||||
# 从检查点恢复训练
|
||||
python train.py --resume checkpoints/mask_rcnn_epoch_20.pth
|
||||
```
|
||||
|
||||
**主要参数说明：**
|
||||
|
||||
| 参数 | 默认值 | 说明 |
|
||||
|------|--------|------|
|
||||
| `--dataset` | `dataset` | 数据集目录路径 |
|
||||
| `--epochs` | `50` | 训练轮数 |
|
||||
| `--batch-size` | `2` | 批大小 |
|
||||
| `--lr` | `0.005` | 学习率 |
|
||||
| `--val-ratio` | `0.2` | 验证集比例 |
|
||||
| `--freeze-backbone` | `False` | 冻结骨干网络 |
|
||||
| `--no-pretrained` | `False` | 不使用预训练权重 |
|
||||
| `--resume` | `None` | 恢复训练的检查点路径 |
|
||||
|
||||
### 2. 推理与可视化
|
||||
|
||||
```bash
|
||||
# 对单张图片推理
|
||||
python predict.py --image dataset/1.jpg --checkpoint checkpoints/mask_rcnn_best.pth
|
||||
|
||||
# 对整个数据集目录批量推理
|
||||
python predict.py --dataset dataset --checkpoint checkpoints/mask_rcnn_best.pth
|
||||
|
||||
# 调整置信度阈值
|
||||
python predict.py --image dataset/1.jpg --threshold 0.7
|
||||
|
||||
# 同时显示 ground truth 对比
|
||||
python predict.py --dataset dataset --show-gt
|
||||
```
|
||||
|
||||
## 数据集格式
|
||||
|
||||
使用 [LabelMe](https://github.com/wkentaro/labelme) 工具进行标注，数据目录结构：
|
||||
|
||||
```
|
||||
dataset/
|
||||
├── image1.jpg
|
||||
├── image1.json
|
||||
├── image2.jpg
|
||||
├── image2.json
|
||||
└── ...
|
||||
```
|
||||
|
||||
每个 JSON 文件的 `shapes` 列表中每个 shape 包含：
|
||||
- `label`: 类别名称（如 `"crosswalk"`）
|
||||
- `points`: `[[x, y], ...]` 坐标点列表
|
||||
- `shape_type`: 标注类型
|
||||
|
||||
### 标注类型说明
|
||||
|
||||
| 类型 | 说明 | 处理方式 |
|
||||
|------|------|----------|
|
||||
| `polygon` | 多边形 | 直接转换为 mask |
|
||||
| `linestrip` | 线段 | 相同标签的线段自动合并为多边形 |
|
||||
| `rectangle` | 矩形 | 转换为四点多边形 |
|
||||
| `circle` | 圆形 | 近似为多边形 |
|
||||
|
||||
> **Linestrip 合并**：当标注对象边界使用多条线段（而非直接画多边形）时，
|
||||
> 代码会自动将相同标签的线段按最近端点连接，合并为闭合多边形。
|
||||
|
||||
## 训练输出
|
||||
|
||||
训练完成后，`checkpoints/` 目录包含：
|
||||
|
||||
| 文件 | 说明 |
|
||||
|------|------|
|
||||
| `mask_rcnn_best.pth` | 验证分数最高的模型 |
|
||||
| `mask_rcnn_final.pth` | 最后一轮的模型 |
|
||||
| `mask_rcnn_epoch_N.pth` | 每 N 轮的检查点 |
|
||||
| `label_map.json` | 标签映射表 |
|
||||
| `loss_history.json` | 损失历史记录 |
|
||||
| `loss_curve.png` | 损失曲线图 |
|
||||
|
||||
## 小数据集建议
|
||||
|
||||
当前数据集仅有 4 张图片，建议：
|
||||
|
||||
1. **使用 `--freeze-backbone`** 冻结骨干网络，大幅减少可训练参数
|
||||
2. **增加 epoch 数**（如 `--epochs 200`），因为样本少需要更多迭代
|
||||
3. **增加数据**：4 张图片不足以训练出泛化能力强的模型，建议扩充到至少 100+ 张
|
||||
4. **使用 polygon 标注**：直接画多边形比 linestrip 更精确
|
||||
|
||||
## 技术细节
|
||||
|
||||
- **模型**：Mask R-CNN + ResNet50-FPN 骨干网络（torchvision 实现）
|
||||
- **优化器**：SGD + momentum(0.9) + weight_decay(0.0005)
|
||||
- **学习率调度**：StepLR，每 10 个 epoch 衰减为 0.1 倍
|
||||
- **数据增强**：水平翻转 + 亮度/对比度抖动（同时变换 mask 和 boxes）
|
||||
+227
@@ -0,0 +1,227 @@
|
||||
"""
|
||||
自定义 Dataset 类
|
||||
支持 LabelMe 标注格式（polygon / linestrip / rectangle / circle）
|
||||
内置数据增强（水平翻转、亮度调整）
|
||||
"""
|
||||
|
||||
import os
|
||||
import json
|
||||
from typing import List, Dict, Tuple, Optional
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
from utils import (
|
||||
get_label_map,
|
||||
shapes_to_masks,
|
||||
get_image_json_pairs,
|
||||
split_dataset,
|
||||
)
|
||||
|
||||
|
||||
class LabelMeDataset(Dataset):
|
||||
"""
|
||||
LabelMe 格式的实例分割数据集。
|
||||
|
||||
数据目录结构:
|
||||
dataset/
|
||||
├── 1.jpg
|
||||
├── 1.json
|
||||
├── 2.jpg
|
||||
├── 2.json
|
||||
└── ...
|
||||
|
||||
每个 JSON 文件包含 LabelMe 标注，shapes 列表中每个 shape 有:
|
||||
- label: 类别名
|
||||
- points: [[x, y], ...] 点坐标
|
||||
- shape_type: "polygon" / "linestrip" / "rectangle" / "circle"
|
||||
|
||||
linestrip 类型会自动将相同 label 的多条线段合并为一个多边形实例。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dataset_dir: str,
|
||||
label_map: Optional[Dict[str, int]] = None,
|
||||
pairs: Optional[List[Tuple[str, str]]] = None,
|
||||
transforms: Optional[object] = None,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
dataset_dir: 数据集目录路径
|
||||
label_map: 标签名→类别ID 映射（若为 None 则自动构建）
|
||||
pairs: 指定使用的 (image_path, json_path) 配对（若为 None 则扫描全部）
|
||||
transforms: 数据增强变换
|
||||
"""
|
||||
self.dataset_dir = dataset_dir
|
||||
self.transforms = transforms
|
||||
|
||||
# 自动构建或使用传入的标签映射
|
||||
if label_map is None:
|
||||
label_map = get_label_map(dataset_dir)
|
||||
self.label_map = label_map
|
||||
|
||||
# 获取图像-标注配对
|
||||
if pairs is None:
|
||||
pairs = get_image_json_pairs(dataset_dir)
|
||||
self.pairs = pairs
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.pairs)
|
||||
|
||||
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, Dict]:
|
||||
image_path, json_path = self.pairs[idx]
|
||||
|
||||
# 读取图像
|
||||
image = Image.open(image_path).convert("RGB")
|
||||
width, height = image.size
|
||||
|
||||
# 读取标注
|
||||
with open(json_path, "r", encoding="utf-8") as f:
|
||||
anno = json.load(f)
|
||||
|
||||
# 确保 imageHeight / imageWidth 与实际一致
|
||||
anno_h = anno.get("imageHeight", height)
|
||||
anno_w = anno.get("imageWidth", width)
|
||||
|
||||
# 转换 shapes → masks, labels, boxes
|
||||
masks, labels, boxes = shapes_to_masks(
|
||||
anno.get("shapes", []),
|
||||
anno_h,
|
||||
anno_w,
|
||||
self.label_map,
|
||||
)
|
||||
|
||||
# 转为 numpy 数组
|
||||
image_np = np.array(image, dtype=np.float32) / 255.0 # [H, W, 3] 0~1
|
||||
|
||||
# 构建 target 字典
|
||||
num_instances = len(labels)
|
||||
target = {
|
||||
"boxes": torch.as_tensor(boxes, dtype=torch.float32),
|
||||
"labels": torch.as_tensor(labels, dtype=torch.int64),
|
||||
"masks": torch.as_tensor(masks, dtype=torch.uint8),
|
||||
"image_id": torch.tensor([idx]),
|
||||
"area": torch.zeros((num_instances,), dtype=torch.float32),
|
||||
"iscrowd": torch.zeros((num_instances,), dtype=torch.int64),
|
||||
}
|
||||
|
||||
# 计算 area
|
||||
if num_instances > 0:
|
||||
target["area"] = (target["masks"] > 0).sum(dim=(1, 2)).float()
|
||||
|
||||
# 应用数据增强
|
||||
if self.transforms is not None:
|
||||
image_np, target = self.transforms(image_np, target)
|
||||
|
||||
# 转换为 tensor [C, H, W]
|
||||
image_tensor = torch.as_tensor(image_np, dtype=torch.float32)
|
||||
if image_tensor.dim() == 3:
|
||||
image_tensor = image_tensor.permute(2, 0, 1)
|
||||
|
||||
return image_tensor, target
|
||||
|
||||
def get_label_map(self) -> Dict[str, int]:
|
||||
return self.label_map
|
||||
|
||||
def get_num_classes(self) -> int:
|
||||
"""返回类别数（不含背景）"""
|
||||
return max(self.label_map.values()) if self.label_map else 0
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 数据增强
|
||||
# ============================================================
|
||||
|
||||
class Compose:
|
||||
"""组合多个变换"""
|
||||
|
||||
def __init__(self, transforms: List):
|
||||
self.transforms = transforms
|
||||
|
||||
def __call__(self, image: np.ndarray, target: dict) -> Tuple[np.ndarray, dict]:
|
||||
for t in self.transforms:
|
||||
image, target = t(image, target)
|
||||
return image, target
|
||||
|
||||
|
||||
class RandomHorizontalFlip:
|
||||
"""随机水平翻转（同时翻转图像、mask、boxes）"""
|
||||
|
||||
def __init__(self, prob: float = 0.5):
|
||||
self.prob = prob
|
||||
|
||||
def __call__(self, image: np.ndarray, target: dict) -> Tuple[np.ndarray, dict]:
|
||||
if np.random.random() < self.prob:
|
||||
image = np.ascontiguousarray(image[:, ::-1, :])
|
||||
_, width = image.shape[:2]
|
||||
|
||||
# 翻转 mask
|
||||
if "masks" in target and len(target["masks"]) > 0:
|
||||
target["masks"] = torch.flip(target["masks"], dims=[2])
|
||||
target["masks"] = target["masks"].contiguous()
|
||||
|
||||
# 翻转 boxes
|
||||
if "boxes" in target and len(target["boxes"]) > 0:
|
||||
boxes = target["boxes"].clone()
|
||||
boxes[:, 0] = width - target["boxes"][:, 2]
|
||||
boxes[:, 2] = width - target["boxes"][:, 0]
|
||||
target["boxes"] = boxes
|
||||
|
||||
return image, target
|
||||
|
||||
|
||||
class RandomBrightness:
|
||||
"""随机亮度调整"""
|
||||
|
||||
def __init__(self, brightness_range: float = 0.2):
|
||||
self.brightness_range = brightness_range
|
||||
|
||||
def __call__(self, image: np.ndarray, target: dict) -> Tuple[np.ndarray, dict]:
|
||||
factor = 1.0 + np.random.uniform(
|
||||
-self.brightness_range, self.brightness_range
|
||||
)
|
||||
image = np.clip(image * factor, 0.0, 1.0)
|
||||
return image, target
|
||||
|
||||
|
||||
class RandomContrast:
|
||||
"""随机对比度调整"""
|
||||
|
||||
def __init__(self, contrast_range: float = 0.2):
|
||||
self.contrast_range = contrast_range
|
||||
|
||||
def __call__(self, image: np.ndarray, target: dict) -> Tuple[np.ndarray, dict]:
|
||||
factor = 1.0 + np.random.uniform(
|
||||
-self.contrast_range, self.contrast_range
|
||||
)
|
||||
mean = image.mean()
|
||||
image = np.clip((image - mean) * factor + mean, 0.0, 1.0)
|
||||
return image, target
|
||||
|
||||
|
||||
class Normalize:
|
||||
"""标准化（使用 ImageNet 均值和标准差）"""
|
||||
|
||||
MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
||||
STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
||||
|
||||
def __call__(self, image: np.ndarray, target: dict) -> Tuple[np.ndarray, dict]:
|
||||
image = (image - self.MEAN) / self.STD
|
||||
return image, target
|
||||
|
||||
|
||||
def get_train_transforms() -> Compose:
|
||||
"""获取训练集数据增强"""
|
||||
return Compose([
|
||||
RandomHorizontalFlip(prob=0.5),
|
||||
RandomBrightness(brightness_range=0.2),
|
||||
RandomContrast(contrast_range=0.2),
|
||||
])
|
||||
|
||||
|
||||
def get_val_transforms() -> Compose:
|
||||
"""获取验证集数据增强（仅标准化，不做随机变换）"""
|
||||
return Compose([])
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 141 KiB |
@@ -0,0 +1,66 @@
|
||||
{
|
||||
"version": "6.3.1",
|
||||
"flags": {},
|
||||
"shapes": [
|
||||
{
|
||||
"label": "crosswalk",
|
||||
"points": [
|
||||
[
|
||||
102.10596026490063,
|
||||
424.4503311258278
|
||||
],
|
||||
[
|
||||
241.17880794701983,
|
||||
254.2516556291391
|
||||
],
|
||||
[
|
||||
355.08609271523176,
|
||||
132.3973509933775
|
||||
],
|
||||
[
|
||||
407.40397350993373,
|
||||
99.94701986754968
|
||||
]
|
||||
],
|
||||
"group_id": null,
|
||||
"description": "",
|
||||
"shape_type": "linestrip",
|
||||
"flags": {},
|
||||
"mask": null
|
||||
},
|
||||
{
|
||||
"label": "crosswalk",
|
||||
"points": [
|
||||
[
|
||||
1006.7417218543045,
|
||||
425.1125827814569
|
||||
],
|
||||
[
|
||||
837.2052980132449,
|
||||
254.91390728476821
|
||||
],
|
||||
[
|
||||
710.7152317880793,
|
||||
147.6291390728477
|
||||
],
|
||||
[
|
||||
657.0728476821191,
|
||||
113.8543046357616
|
||||
],
|
||||
[
|
||||
625.2847682119204,
|
||||
100.60927152317882
|
||||
]
|
||||
],
|
||||
"group_id": null,
|
||||
"description": "",
|
||||
"shape_type": "linestrip",
|
||||
"flags": {},
|
||||
"mask": null
|
||||
}
|
||||
],
|
||||
"imagePath": "1.jpg",
|
||||
"imageData": null,
|
||||
"imageHeight": 584,
|
||||
"imageWidth": 1036
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 102 KiB |
@@ -0,0 +1,46 @@
|
||||
{
|
||||
"version": "6.3.1",
|
||||
"flags": {},
|
||||
"shapes": [
|
||||
{
|
||||
"label": "crosswalk",
|
||||
"points": [
|
||||
[
|
||||
33.42372881355931,
|
||||
75.98870056497175
|
||||
],
|
||||
[
|
||||
589.3559322033898,
|
||||
77.11864406779661
|
||||
]
|
||||
],
|
||||
"group_id": null,
|
||||
"description": "",
|
||||
"shape_type": "linestrip",
|
||||
"flags": {},
|
||||
"mask": null
|
||||
},
|
||||
{
|
||||
"label": "crosswalk",
|
||||
"points": [
|
||||
[
|
||||
587.0960451977401,
|
||||
447.7401129943502
|
||||
],
|
||||
[
|
||||
26.64406779661016,
|
||||
448.8700564971751
|
||||
]
|
||||
],
|
||||
"group_id": null,
|
||||
"description": "",
|
||||
"shape_type": "linestrip",
|
||||
"flags": {},
|
||||
"mask": null
|
||||
}
|
||||
],
|
||||
"imagePath": "2.jpg",
|
||||
"imageData": null,
|
||||
"imageHeight": 500,
|
||||
"imageWidth": 616
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 64 KiB |
@@ -0,0 +1,46 @@
|
||||
{
|
||||
"version": "6.3.1",
|
||||
"flags": {},
|
||||
"shapes": [
|
||||
{
|
||||
"label": "crosswalk",
|
||||
"points": [
|
||||
[
|
||||
0.0,
|
||||
0.42372881355932207
|
||||
],
|
||||
[
|
||||
0.0,
|
||||
681.7796610169491
|
||||
]
|
||||
],
|
||||
"group_id": null,
|
||||
"description": "",
|
||||
"shape_type": "linestrip",
|
||||
"flags": {},
|
||||
"mask": null
|
||||
},
|
||||
{
|
||||
"label": "crosswalk",
|
||||
"points": [
|
||||
[
|
||||
499.1525423728813,
|
||||
671.6101694915254
|
||||
],
|
||||
[
|
||||
498.3050847457627,
|
||||
2.1186440677966103
|
||||
]
|
||||
],
|
||||
"group_id": null,
|
||||
"description": "",
|
||||
"shape_type": "linestrip",
|
||||
"flags": {},
|
||||
"mask": null
|
||||
}
|
||||
],
|
||||
"imagePath": "3.jpg",
|
||||
"imageData": null,
|
||||
"imageHeight": 750,
|
||||
"imageWidth": 500
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 128 KiB |
@@ -0,0 +1,46 @@
|
||||
{
|
||||
"version": "6.3.1",
|
||||
"flags": {},
|
||||
"shapes": [
|
||||
{
|
||||
"label": "crosswalk",
|
||||
"points": [
|
||||
[
|
||||
1.1187050359712742,
|
||||
434.9856115107914
|
||||
],
|
||||
[
|
||||
445.00359712230215,
|
||||
234.26618705035972
|
||||
]
|
||||
],
|
||||
"group_id": null,
|
||||
"description": "",
|
||||
"shape_type": "linestrip",
|
||||
"flags": {},
|
||||
"mask": null
|
||||
},
|
||||
{
|
||||
"label": "crosswalk",
|
||||
"points": [
|
||||
[
|
||||
947.8812949640287,
|
||||
539.3021582733813
|
||||
],
|
||||
[
|
||||
588.1690647482014,
|
||||
228.5107913669065
|
||||
]
|
||||
],
|
||||
"group_id": null,
|
||||
"description": "",
|
||||
"shape_type": "linestrip",
|
||||
"flags": {},
|
||||
"mask": null
|
||||
}
|
||||
],
|
||||
"imagePath": "4.jpg",
|
||||
"imageData": null,
|
||||
"imageHeight": 634,
|
||||
"imageWidth": 949
|
||||
}
|
||||
@@ -0,0 +1,170 @@
|
||||
"""
|
||||
Mask R-CNN 模型构建模块
|
||||
基于 torchvision 的 maskrcnn_resnet50_fpn，支持预训练权重微调
|
||||
优先从本地路径加载 COCO 预训练权重，避免联网下载
|
||||
"""
|
||||
|
||||
import os
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torchvision
|
||||
from torchvision.models.detection import (
|
||||
maskrcnn_resnet50_fpn,
|
||||
MaskRCNN_ResNet50_FPN_Weights,
|
||||
)
|
||||
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
|
||||
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
|
||||
|
||||
|
||||
# 本地 COCO 预训练权重路径
|
||||
_LOCAL_COCO_WEIGHTS = os.path.join(
|
||||
os.path.dirname(os.path.abspath(__file__)),
|
||||
"model",
|
||||
"maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth",
|
||||
)
|
||||
|
||||
|
||||
def get_model(
|
||||
num_classes: int,
|
||||
pretrained: bool = True,
|
||||
pretrained_backbone: bool = True,
|
||||
min_size: int = 800,
|
||||
max_size: int = 1333,
|
||||
pretrained_path: str = None,
|
||||
) -> nn.Module:
|
||||
"""
|
||||
构建 Mask R-CNN 模型（ResNet50-FPN 骨干网络）。
|
||||
|
||||
如果 pretrained=True，加载 COCO 预训练权重，仅替换最后的分类头和 mask 预测头
|
||||
以适配自定义类别数。这大幅加快收敛速度，特别适合小数据集。
|
||||
|
||||
权重加载优先级：
|
||||
1. pretrained_path 指定的路径
|
||||
2. 本地 model/maskrcnn_resnet50_fpn_coco-bf2d0c1e.pth
|
||||
3. 在线下载（ torchvision 默认行为）
|
||||
|
||||
Args:
|
||||
num_classes: 类别数（包含背景，如 2 类 = 背景 + 1 个前景类）
|
||||
pretrained: 是否加载 COCO 预训练权重
|
||||
pretrained_backbone: 是否加载 ImageNet 预训练骨干网络权重
|
||||
min_size: 输入图像最小边 resize 后的尺寸
|
||||
max_size: 输入图像最大边 resize 后的尺寸
|
||||
pretrained_path: 自定义预训练权重路径（优先于本地默认路径）
|
||||
|
||||
Returns:
|
||||
Mask R-CNN 模型
|
||||
"""
|
||||
# 确定本地权重路径
|
||||
local_path = pretrained_path or _LOCAL_COCO_WEIGHTS
|
||||
has_local = os.path.isfile(local_path)
|
||||
|
||||
if pretrained:
|
||||
if has_local:
|
||||
# 优先从本地加载权重，不触发在线下载
|
||||
print(f"从本地加载 COCO 预训练权重: {local_path}")
|
||||
model = maskrcnn_resnet50_fpn(
|
||||
weights=None,
|
||||
weights_backbone=None,
|
||||
min_size=min_size,
|
||||
max_size=max_size,
|
||||
)
|
||||
state_dict = torch.load(local_path, map_location="cpu", weights_only=True)
|
||||
model.load_state_dict(state_dict)
|
||||
print(" -> 本地权重加载成功")
|
||||
else:
|
||||
# 本地文件不存在，回退到在线下载
|
||||
print(f"本地权重不存在 ({local_path})，从在线下载 COCO 预训练权重...")
|
||||
model = maskrcnn_resnet50_fpn(
|
||||
weights=MaskRCNN_ResNet50_FPN_Weights.COCO_V1,
|
||||
weights_backbone=None,
|
||||
min_size=min_size,
|
||||
max_size=max_size,
|
||||
)
|
||||
else:
|
||||
model = maskrcnn_resnet50_fpn(
|
||||
weights=None,
|
||||
weights_backbone=(
|
||||
torchvision.models.ResNet50_Weights.IMAGENET1K_V1
|
||||
if pretrained_backbone else None
|
||||
),
|
||||
min_size=min_size,
|
||||
max_size=max_size,
|
||||
)
|
||||
|
||||
# 替换分类头（box predictor）
|
||||
in_features = model.roi_heads.box_predictor.cls_score.in_features
|
||||
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
|
||||
|
||||
# 替换 mask 预测头
|
||||
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
|
||||
hidden_layer = 256
|
||||
model.roi_heads.mask_predictor = MaskRCNNPredictor(
|
||||
in_features_mask, hidden_layer, num_classes
|
||||
)
|
||||
|
||||
return model
|
||||
|
||||
|
||||
def freeze_backbone(model: nn.Module, freeze: bool = True) -> nn.Module:
|
||||
"""
|
||||
冻结/解冻骨干网络参数。
|
||||
冻结骨干网络可以减少训练参数量，适合小数据集场景。
|
||||
|
||||
Args:
|
||||
model: Mask R-CNN 模型
|
||||
freeze: 是否冻结
|
||||
|
||||
Returns:
|
||||
修改后的模型
|
||||
"""
|
||||
for param in model.backbone.parameters():
|
||||
param.requires_grad = not freeze
|
||||
return model
|
||||
|
||||
|
||||
def get_optimizer(
|
||||
model: nn.Module,
|
||||
lr: float = 0.005,
|
||||
momentum: float = 0.9,
|
||||
weight_decay: float = 0.0005,
|
||||
) -> torch.optim.Optimizer:
|
||||
"""
|
||||
创建 SGD 优化器（Mask R-CNN 的标准配置）。
|
||||
可学习参数会被分组，冻结的参数不包含在优化器中。
|
||||
|
||||
Args:
|
||||
model: 模型
|
||||
lr: 学习率
|
||||
momentum: SGD 动量
|
||||
weight_decay: 权重衰减（L2 正则化）
|
||||
|
||||
Returns:
|
||||
SGD 优化器
|
||||
"""
|
||||
params = [p for p in model.parameters() if p.requires_grad]
|
||||
optimizer = torch.optim.SGD(
|
||||
params,
|
||||
lr=lr,
|
||||
momentum=momentum,
|
||||
weight_decay=weight_decay,
|
||||
)
|
||||
return optimizer
|
||||
|
||||
|
||||
def get_lr_scheduler(
|
||||
optimizer: torch.optim.Optimizer,
|
||||
step_size: int = 5,
|
||||
gamma: float = 0.1,
|
||||
) -> torch.optim.lr_scheduler.StepLR:
|
||||
"""
|
||||
创建学习率调度器（StepLR，每 step_size 个 epoch 衰减为 gamma 倍）。
|
||||
|
||||
Args:
|
||||
optimizer: 优化器
|
||||
step_size: 衰减间隔（epoch 数）
|
||||
gamma: 衰减系数
|
||||
|
||||
Returns:
|
||||
StepLR 调度器
|
||||
"""
|
||||
return torch.optim.lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=gamma)
|
||||
+262
@@ -0,0 +1,262 @@
|
||||
"""
|
||||
Mask R-CNN 实例分割推理与可视化脚本
|
||||
|
||||
用法:
|
||||
# 对单张图片进行推理
|
||||
python predict.py --image dataset/1.jpg --checkpoint checkpoints/mask_rcnn_best.pth
|
||||
|
||||
# 对整个数据集目录进行推理
|
||||
python predict.py --dataset dataset --checkpoint checkpoints/mask_rcnn_best.pth
|
||||
|
||||
# 调整置信度阈值
|
||||
python predict.py --image dataset/1.jpg --checkpoint checkpoints/mask_rcnn_best.pth --threshold 0.7
|
||||
|
||||
# 同时显示 ground truth 对比
|
||||
python predict.py --image dataset/1.jpg --checkpoint checkpoints/mask_rcnn_best.pth --show-gt
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
from typing import Dict, Optional
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import torch
|
||||
|
||||
from utils import (
|
||||
get_label_map,
|
||||
get_image_json_pairs,
|
||||
shapes_to_masks,
|
||||
visualize_instances,
|
||||
)
|
||||
from dataset import get_val_transforms
|
||||
from model import get_model
|
||||
|
||||
|
||||
def load_checkpoint(checkpoint_path: str, device: torch.device):
|
||||
"""
|
||||
加载训练好的模型检查点。
|
||||
|
||||
Args:
|
||||
checkpoint_path: 检查点文件路径
|
||||
device: 计算设备
|
||||
|
||||
Returns:
|
||||
model: 加载好权重的模型
|
||||
label_map: 标签名→类别ID 映射
|
||||
"""
|
||||
checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
|
||||
label_map = checkpoint.get("label_map", {})
|
||||
|
||||
num_classes = max(label_map.values()) + 1 # +1 for background
|
||||
|
||||
model = get_model(num_classes=num_classes, pretrained=False)
|
||||
model.load_state_dict(checkpoint["model_state_dict"])
|
||||
model = model.to(device)
|
||||
model.eval()
|
||||
|
||||
# 构建 id→name 映射
|
||||
id_to_name = {v: k for k, v in label_map.items()}
|
||||
|
||||
return model, label_map, id_to_name
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def predict_single(
|
||||
model: torch.nn.Module,
|
||||
image: np.ndarray,
|
||||
device: torch.device,
|
||||
threshold: float = 0.5,
|
||||
) -> Dict:
|
||||
"""
|
||||
对单张图像进行推理。
|
||||
|
||||
Args:
|
||||
model: 模型
|
||||
image: [H, W, 3] RGB uint8 图像
|
||||
device: 计算设备
|
||||
threshold: 置信度阈值
|
||||
|
||||
Returns:
|
||||
包含 boxes, masks, labels, scores 的字典
|
||||
"""
|
||||
# 预处理
|
||||
image_float = image.astype(np.float32) / 255.0
|
||||
image_tensor = torch.as_tensor(image_float, dtype=torch.float32)
|
||||
image_tensor = image_tensor.permute(2, 0, 1).unsqueeze(0).to(device)
|
||||
|
||||
# 推理
|
||||
outputs = model(image_tensor)
|
||||
output = outputs[0]
|
||||
|
||||
# 过滤低置信度
|
||||
keep = output["scores"] > threshold
|
||||
boxes = output["boxes"][keep].cpu().numpy()
|
||||
masks = output["masks"][keep].cpu().numpy()
|
||||
labels = output["labels"][keep].cpu().numpy()
|
||||
scores = output["scores"][keep].cpu().numpy()
|
||||
|
||||
# mask 阈值化（> 0.5）
|
||||
masks = (masks[:, 0, :, :] > 0.5).astype(np.uint8)
|
||||
|
||||
return {
|
||||
"boxes": boxes,
|
||||
"masks": masks,
|
||||
"labels": labels,
|
||||
"scores": scores,
|
||||
}
|
||||
|
||||
|
||||
def visualize_and_save(
|
||||
image: np.ndarray,
|
||||
result: Dict,
|
||||
id_to_name: Dict[int, str],
|
||||
output_path: str,
|
||||
threshold: float = 0.5,
|
||||
):
|
||||
"""
|
||||
可视化预测结果并保存图像。
|
||||
|
||||
Args:
|
||||
image: 原始图像 [H, W, 3] RGB
|
||||
result: 预测结果字典
|
||||
id_to_name: 类别ID→名称映射
|
||||
output_path: 输出路径
|
||||
threshold: 置信度阈值
|
||||
"""
|
||||
vis_img = visualize_instances(
|
||||
image,
|
||||
result["boxes"],
|
||||
result["masks"],
|
||||
result["labels"],
|
||||
result["scores"],
|
||||
label_map=id_to_name,
|
||||
score_threshold=threshold,
|
||||
)
|
||||
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
Image.fromarray(vis_img).save(output_path)
|
||||
print(f" -> 结果已保存: {output_path}")
|
||||
|
||||
|
||||
def visualize_gt(
|
||||
image: np.ndarray,
|
||||
json_path: str,
|
||||
label_map: Dict[str, int],
|
||||
id_to_name: Dict[int, str],
|
||||
output_path: str,
|
||||
):
|
||||
"""可视化 ground truth 标注"""
|
||||
with open(json_path, "r", encoding="utf-8") as f:
|
||||
anno = json.load(f)
|
||||
|
||||
height = anno.get("imageHeight", image.shape[0])
|
||||
width = anno.get("imageWidth", image.shape[1])
|
||||
|
||||
masks, labels, boxes = shapes_to_masks(
|
||||
anno.get("shapes", []), height, width, label_map
|
||||
)
|
||||
|
||||
if len(labels) == 0:
|
||||
print(" -> 无标注实例")
|
||||
return
|
||||
|
||||
vis_img = visualize_instances(
|
||||
image, boxes, masks, labels, None, id_to_name, 0.0
|
||||
)
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
Image.fromarray(vis_img).save(output_path)
|
||||
print(f" -> Ground Truth 已保存: {output_path}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Mask R-CNN 实例分割推理")
|
||||
parser.add_argument("--image", type=str, default=None,
|
||||
help="单张图像路径")
|
||||
parser.add_argument("--dataset", type=str, default=None,
|
||||
help="数据集目录路径（批量推理）")
|
||||
parser.add_argument("--checkpoint", type=str, default="checkpoints/mask_rcnn_best.pth",
|
||||
help="模型检查点路径 (默认: checkpoints/mask_rcnn_best.pth)")
|
||||
parser.add_argument("--output-dir", type=str, default="predictions",
|
||||
help="输出目录 (默认: predictions)")
|
||||
parser.add_argument("--threshold", type=float, default=0.5,
|
||||
help="置信度阈值 (默认: 0.5)")
|
||||
parser.add_argument("--show-gt", action="store_true",
|
||||
help="同时可视化 ground truth 对比")
|
||||
parser.add_argument("--min-size", type=int, default=800,
|
||||
help="输入图像最小边尺寸 (默认: 800)")
|
||||
parser.add_argument("--max-size", type=int, default=1333,
|
||||
help="输入图像最大边尺寸 (默认: 1333)")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# 设置设备
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
print(f"使用设备: {device}")
|
||||
|
||||
# 检查检查点
|
||||
ckpt_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), args.checkpoint)
|
||||
if not os.path.exists(ckpt_path):
|
||||
print(f"错误: 检查点文件不存在: {ckpt_path}")
|
||||
print("请先运行 train.py 训练模型，或使用 --checkpoint 指定正确的路径。")
|
||||
sys.exit(1)
|
||||
|
||||
# 加载模型
|
||||
model, label_map, id_to_name = load_checkpoint(ckpt_path, device)
|
||||
print(f"模型已加载 | 标签映射: {label_map}")
|
||||
print(f"置信度阈值: {args.threshold}")
|
||||
|
||||
output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), args.output_dir)
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# 收集要推理的图像
|
||||
if args.image:
|
||||
image_paths = [args.image]
|
||||
json_paths = [args.image.rsplit(".", 1)[0] + ".json"]
|
||||
elif args.dataset:
|
||||
dataset_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), args.dataset)
|
||||
pairs = get_image_json_pairs(dataset_path)
|
||||
image_paths = [p[0] for p in pairs]
|
||||
json_paths = [p[1] for p in pairs]
|
||||
else:
|
||||
print("错误: 请指定 --image 或 --dataset")
|
||||
sys.exit(1)
|
||||
|
||||
print(f"\n共 {len(image_paths)} 张图像待推理\n")
|
||||
|
||||
for idx, (image_path, json_path) in enumerate(zip(image_paths, json_paths)):
|
||||
if not os.path.exists(image_path):
|
||||
print(f"[{idx+1}] 跳过（图像不存在）: {image_path}")
|
||||
continue
|
||||
|
||||
print(f"[{idx+1}/{len(image_paths)}] {os.path.basename(image_path)}")
|
||||
|
||||
# 读取图像
|
||||
image = np.array(Image.open(image_path).convert("RGB"))
|
||||
|
||||
# 推理
|
||||
result = predict_single(model, image, device, args.threshold)
|
||||
|
||||
num_detections = len(result["scores"])
|
||||
print(f" -> 检测到 {num_detections} 个实例")
|
||||
for i in range(num_detections):
|
||||
name = id_to_name.get(int(result["labels"][i]), str(result["labels"][i]))
|
||||
print(f" {name}: {result['scores'][i]:.3f}")
|
||||
|
||||
# 可视化并保存
|
||||
base_name = os.path.splitext(os.path.basename(image_path))[0]
|
||||
output_path = os.path.join(output_dir, f"{base_name}_pred.png")
|
||||
visualize_and_save(image, result, id_to_name, output_path, args.threshold)
|
||||
|
||||
# 可选：可视化 ground truth
|
||||
if args.show_gt and os.path.exists(json_path):
|
||||
gt_output_path = os.path.join(output_dir, f"{base_name}_gt.png")
|
||||
visualize_gt(image, json_path, label_map, id_to_name, gt_output_path)
|
||||
|
||||
print(f"\n推理完成！结果保存在: {output_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,19 @@
|
||||
# Mask R-CNN 实例分割训练依赖
|
||||
# Python 3.11
|
||||
|
||||
# PyTorch (CPU 版本, 无 GPU 时使用)
|
||||
# 安装命令: pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
|
||||
torch>=2.0.0
|
||||
torchvision>=0.15.0
|
||||
|
||||
# 图像处理
|
||||
opencv-python>=4.8.0
|
||||
Pillow>=9.0.0
|
||||
numpy>=1.24.0
|
||||
scikit-image>=0.21.0
|
||||
|
||||
# 标注工具
|
||||
labelme>=6.0.0
|
||||
|
||||
# 进度条
|
||||
tqdm>=4.65.0
|
||||
@@ -0,0 +1,481 @@
|
||||
"""
|
||||
Mask R-CNN 实例分割训练脚本
|
||||
|
||||
用法:
|
||||
# 使用默认参数训练
|
||||
python train.py
|
||||
|
||||
# 自定义参数训练
|
||||
python train.py --dataset dataset --epochs 50 --batch-size 2 --lr 0.005
|
||||
|
||||
# 冻结骨干网络（小数据集推荐）
|
||||
python train.py --freeze-backbone --epochs 100
|
||||
|
||||
# 从检查点恢复训练
|
||||
python train.py --resume checkpoints/mask_rcnn_epoch_20.pth
|
||||
|
||||
功能:
|
||||
1. 自动扫描 dataset 目录中的 LabelMe 标注
|
||||
2. 支持 polygon / linestrip / rectangle / circle 标注类型
|
||||
3. 自动划分训练集/验证集
|
||||
4. 支持 COCO 预训练权重微调
|
||||
5. 支持冻结骨干网络（减少过拟合）
|
||||
6. 训练过程中保存检查点和最佳模型
|
||||
7. 记录训练日志（损失曲线）
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import json
|
||||
from typing import Dict, List
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from utils import get_label_map, get_image_json_pairs, split_dataset, visualize_instances
|
||||
from dataset import LabelMeDataset, get_train_transforms, get_val_transforms
|
||||
from model import get_model, freeze_backbone, get_optimizer, get_lr_scheduler
|
||||
|
||||
|
||||
def collate_fn(batch):
|
||||
"""
|
||||
自定义 collate 函数：Mask R-CNN 的每个样本可能有不同数量的实例，
|
||||
不能使用默认的 stack 方式，需要返回 list。
|
||||
"""
|
||||
return tuple(zip(*batch))
|
||||
|
||||
|
||||
def train_one_epoch(
|
||||
model: torch.nn.Module,
|
||||
optimizer: torch.optim.Optimizer,
|
||||
data_loader: DataLoader,
|
||||
device: torch.device,
|
||||
epoch: int,
|
||||
log_interval: int = 10,
|
||||
) -> Dict[str, float]:
|
||||
"""
|
||||
训练一个 epoch。
|
||||
|
||||
Args:
|
||||
model: 模型
|
||||
optimizer: 优化器
|
||||
data_loader: 训练数据加载器
|
||||
device: 计算设备 (cuda / cpu)
|
||||
epoch: 当前 epoch 编号
|
||||
log_interval: 日志打印间隔
|
||||
|
||||
Returns:
|
||||
各项损失的平均值
|
||||
"""
|
||||
model.train()
|
||||
total_loss = 0.0
|
||||
loss_components = {
|
||||
"loss_classifier": 0.0,
|
||||
"loss_box_reg": 0.0,
|
||||
"loss_mask": 0.0,
|
||||
"loss_objectness": 0.0,
|
||||
"loss_rpn_box_reg": 0.0,
|
||||
}
|
||||
num_batches = 0
|
||||
|
||||
try:
|
||||
from tqdm import tqdm
|
||||
pbar = tqdm(data_loader, desc=f"Epoch {epoch}")
|
||||
except ImportError:
|
||||
pbar = data_loader
|
||||
|
||||
for batch_idx, (images, targets) in enumerate(pbar):
|
||||
# 过滤无实例的样本
|
||||
valid_images = []
|
||||
valid_targets = []
|
||||
for img, tgt in zip(images, targets):
|
||||
if len(tgt["boxes"]) > 0:
|
||||
valid_images.append(img)
|
||||
valid_targets.append(tgt)
|
||||
|
||||
if len(valid_images) == 0:
|
||||
continue
|
||||
|
||||
images = [img.to(device) for img in valid_images]
|
||||
targets = [{k: v.to(device) for k, v in t.items()} for t in valid_targets]
|
||||
|
||||
# 前向传播（训练模式下返回 loss dict）
|
||||
loss_dict = model(images, targets)
|
||||
losses = sum(loss for loss in loss_dict.values())
|
||||
|
||||
# 反向传播
|
||||
optimizer.zero_grad()
|
||||
losses.backward()
|
||||
optimizer.step()
|
||||
|
||||
# 记录损失
|
||||
total_loss += losses.item()
|
||||
num_batches += 1
|
||||
for k, v in loss_dict.items():
|
||||
loss_components[k] += v.item()
|
||||
|
||||
# 更新进度条
|
||||
if hasattr(pbar, 'set_postfix'):
|
||||
pbar.set_postfix({"loss": f"{losses.item():.4f}"})
|
||||
|
||||
avg_loss = total_loss / max(num_batches, 1)
|
||||
for k in loss_components:
|
||||
loss_components[k] /= max(num_batches, 1)
|
||||
loss_components["total"] = avg_loss
|
||||
|
||||
return loss_components
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def evaluate(
|
||||
model: torch.nn.Module,
|
||||
data_loader: DataLoader,
|
||||
device: torch.device,
|
||||
) -> float:
|
||||
"""
|
||||
验证集评估：计算平均检测置信度作为近似指标。
|
||||
（完整的 COCO mAP 评估需要 pycocotools，这里使用简化指标避免额外依赖）
|
||||
|
||||
Args:
|
||||
model: 模型
|
||||
data_loader: 验证数据加载器
|
||||
device: 计算设备
|
||||
|
||||
Returns:
|
||||
平均置信度（0~1）
|
||||
"""
|
||||
model.eval()
|
||||
total_score = 0.0
|
||||
total_instances = 0
|
||||
|
||||
for images, targets in data_loader:
|
||||
images = [img.to(device) for img in images]
|
||||
outputs = model(images)
|
||||
|
||||
for output in outputs:
|
||||
scores = output["scores"]
|
||||
if len(scores) > 0:
|
||||
# 取置信度 > 0.5 的预测
|
||||
high_conf = scores[scores > 0.5]
|
||||
if len(high_conf) > 0:
|
||||
total_score += high_conf.mean().item()
|
||||
total_instances += 1
|
||||
|
||||
return total_score / max(total_instances, 1)
|
||||
|
||||
|
||||
def save_checkpoint(
|
||||
model: torch.nn.Module,
|
||||
optimizer: torch.optim.Optimizer,
|
||||
scheduler,
|
||||
epoch: int,
|
||||
label_map: Dict[str, int],
|
||||
loss_history: List[Dict],
|
||||
path: str,
|
||||
):
|
||||
"""保存训练检查点"""
|
||||
torch.save({
|
||||
"epoch": epoch,
|
||||
"model_state_dict": model.state_dict(),
|
||||
"optimizer_state_dict": optimizer.state_dict(),
|
||||
"scheduler_state_dict": scheduler.state_dict() if scheduler else None,
|
||||
"label_map": label_map,
|
||||
"loss_history": loss_history,
|
||||
}, path)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Mask R-CNN 实例分割训练")
|
||||
parser.add_argument("--dataset", type=str, default="dataset",
|
||||
help="数据集目录路径 (默认: dataset)")
|
||||
parser.add_argument("--output-dir", type=str, default="checkpoints",
|
||||
help="检查点输出目录 (默认: checkpoints)")
|
||||
parser.add_argument("--epochs", type=int, default=50,
|
||||
help="训练轮数 (默认: 50)")
|
||||
parser.add_argument("--batch-size", type=int, default=2,
|
||||
help="批大小 (默认: 2)")
|
||||
parser.add_argument("--lr", type=float, default=0.005,
|
||||
help="学习率 (默认: 0.005)")
|
||||
parser.add_argument("--momentum", type=float, default=0.9,
|
||||
help="SGD 动量 (默认: 0.9)")
|
||||
parser.add_argument("--weight-decay", type=float, default=0.0005,
|
||||
help="权重衰减 (默认: 0.0005)")
|
||||
parser.add_argument("--lr-step-size", type=int, default=10,
|
||||
help="学习率衰减间隔 (默认: 10)")
|
||||
parser.add_argument("--lr-gamma", type=float, default=0.1,
|
||||
help="学习率衰减系数 (默认: 0.1)")
|
||||
parser.add_argument("--val-ratio", type=float, default=0.2,
|
||||
help="验证集比例 (默认: 0.2)")
|
||||
parser.add_argument("--num-workers", type=int, default=0,
|
||||
help="数据加载线程数 (Windows 建议 0) (默认: 0)")
|
||||
parser.add_argument("--no-pretrained", action="store_true",
|
||||
help="不使用 COCO 预训练权重")
|
||||
parser.add_argument("--freeze-backbone", action="store_true",
|
||||
help="冻结骨干网络（小数据集推荐）")
|
||||
parser.add_argument("--resume", type=str, default=None,
|
||||
help="从检查点恢复训练的路径")
|
||||
parser.add_argument("--save-freq", type=int, default=5,
|
||||
help="每多少 epoch 保存一次检查点 (默认: 5)")
|
||||
parser.add_argument("--min-size", type=int, default=800,
|
||||
help="输入图像最小边尺寸 (默认: 800)")
|
||||
parser.add_argument("--max-size", type=int, default=1333,
|
||||
help="输入图像最大边尺寸 (默认: 1333)")
|
||||
parser.add_argument("--seed", type=int, default=42,
|
||||
help="随机种子 (默认: 42)")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# 设置随机种子
|
||||
torch.manual_seed(args.seed)
|
||||
np.random.seed(args.seed)
|
||||
|
||||
# 设置设备
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
print(f"使用设备: {device}")
|
||||
|
||||
# 构建标签映射
|
||||
dataset_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), args.dataset)
|
||||
label_map = get_label_map(dataset_path)
|
||||
print(f"标签映射: {label_map}")
|
||||
|
||||
num_classes = max(label_map.values()) + 1 # +1 for background
|
||||
print(f"类别数（含背景）: {num_classes}")
|
||||
|
||||
# 获取数据配对并划分
|
||||
all_pairs = get_image_json_pairs(dataset_path)
|
||||
print(f"数据集总样本数: {len(all_pairs)}")
|
||||
|
||||
train_pairs, val_pairs = split_dataset(all_pairs, val_ratio=args.val_ratio, seed=args.seed)
|
||||
print(f"训练集: {len(train_pairs)} | 验证集: {len(val_pairs)}")
|
||||
|
||||
if len(train_pairs) == 0:
|
||||
print("错误: 训练集为空！请检查数据集目录。")
|
||||
sys.exit(1)
|
||||
|
||||
# 创建数据集和数据加载器
|
||||
train_dataset = LabelMeDataset(
|
||||
dataset_dir=dataset_path,
|
||||
label_map=label_map,
|
||||
pairs=train_pairs,
|
||||
transforms=get_train_transforms(),
|
||||
)
|
||||
val_dataset = LabelMeDataset(
|
||||
dataset_dir=dataset_path,
|
||||
label_map=label_map,
|
||||
pairs=val_pairs,
|
||||
transforms=get_val_transforms(),
|
||||
) if len(val_pairs) > 0 else None
|
||||
|
||||
train_loader = DataLoader(
|
||||
train_dataset,
|
||||
batch_size=args.batch_size,
|
||||
shuffle=True,
|
||||
num_workers=args.num_workers,
|
||||
collate_fn=collate_fn,
|
||||
drop_last=False,
|
||||
)
|
||||
val_loader = DataLoader(
|
||||
val_dataset,
|
||||
batch_size=1,
|
||||
shuffle=False,
|
||||
num_workers=args.num_workers,
|
||||
collate_fn=collate_fn,
|
||||
) if val_dataset else None
|
||||
|
||||
# 创建模型
|
||||
pretrained = not args.no_pretrained
|
||||
model = get_model(
|
||||
num_classes=num_classes,
|
||||
pretrained=pretrained,
|
||||
min_size=args.min_size,
|
||||
max_size=args.max_size,
|
||||
)
|
||||
print(f"预训练权重: {'COCO' if pretrained else '无'}")
|
||||
|
||||
if args.freeze_backbone:
|
||||
model = freeze_backbone(model, freeze=True)
|
||||
num_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||||
num_total = sum(p.numel() for p in model.parameters())
|
||||
print(f"骨干网络已冻结 | 可训练参数: {num_trainable:,} / {num_total:,}")
|
||||
else:
|
||||
num_total = sum(p.numel() for p in model.parameters())
|
||||
print(f"可训练参数: {num_total:,}")
|
||||
|
||||
model = model.to(device)
|
||||
|
||||
# 创建优化器和调度器
|
||||
optimizer = get_optimizer(model, lr=args.lr, momentum=args.momentum,
|
||||
weight_decay=args.weight_decay)
|
||||
scheduler = get_lr_scheduler(optimizer, step_size=args.lr_step_size, gamma=args.lr_gamma)
|
||||
|
||||
# 恢复训练
|
||||
start_epoch = 0
|
||||
loss_history = []
|
||||
best_score = 0.0
|
||||
|
||||
if args.resume and os.path.exists(args.resume):
|
||||
checkpoint = torch.load(args.resume, map_location=device, weights_only=False)
|
||||
model.load_state_dict(checkpoint["model_state_dict"])
|
||||
optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
|
||||
if checkpoint.get("scheduler_state_dict") and scheduler:
|
||||
scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
|
||||
start_epoch = checkpoint["epoch"] + 1
|
||||
loss_history = checkpoint.get("loss_history", [])
|
||||
best_score = max((h.get("val_score", 0) for h in loss_history), default=0.0)
|
||||
print(f"从 epoch {start_epoch} 恢复训练 (best_score={best_score:.4f})")
|
||||
|
||||
# 创建输出目录
|
||||
output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), args.output_dir)
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# 保存标签映射
|
||||
with open(os.path.join(output_dir, "label_map.json"), "w", encoding="utf-8") as f:
|
||||
json.dump(label_map, f, ensure_ascii=False, indent=2)
|
||||
|
||||
# 训练循环
|
||||
print(f"\n{'='*60}")
|
||||
print(f"开始训练 | 总轮数: {args.epochs} | 批大小: {args.batch_size} | 学习率: {args.lr}")
|
||||
print(f"{'='*60}\n")
|
||||
|
||||
for epoch in range(start_epoch, args.epochs):
|
||||
epoch_start = time.time()
|
||||
|
||||
# 训练
|
||||
train_metrics = train_one_epoch(
|
||||
model, optimizer, train_loader, device, epoch + 1
|
||||
)
|
||||
|
||||
# 验证
|
||||
val_score = 0.0
|
||||
if val_loader is not None:
|
||||
val_score = evaluate(model, val_loader, device)
|
||||
|
||||
# 更新学习率
|
||||
scheduler.step()
|
||||
|
||||
epoch_time = time.time() - epoch_start
|
||||
current_lr = optimizer.param_groups[0]["lr"]
|
||||
|
||||
# 记录损失历史
|
||||
epoch_record = {
|
||||
"epoch": epoch + 1,
|
||||
"train_loss": train_metrics["total"],
|
||||
"loss_classifier": train_metrics["loss_classifier"],
|
||||
"loss_box_reg": train_metrics["loss_box_reg"],
|
||||
"loss_mask": train_metrics["loss_mask"],
|
||||
"loss_objectness": train_metrics["loss_objectness"],
|
||||
"loss_rpn_box_reg": train_metrics["loss_rpn_box_reg"],
|
||||
"val_score": val_score,
|
||||
"lr": current_lr,
|
||||
"time": round(epoch_time, 1),
|
||||
}
|
||||
loss_history.append(epoch_record)
|
||||
|
||||
# 打印日志
|
||||
print(
|
||||
f"Epoch {epoch+1}/{args.epochs} | "
|
||||
f"Loss: {train_metrics['total']:.4f} "
|
||||
f"(cls:{train_metrics['loss_classifier']:.4f} "
|
||||
f"box:{train_metrics['loss_box_reg']:.4f} "
|
||||
f"mask:{train_metrics['loss_mask']:.4f}) | "
|
||||
f"Val: {val_score:.4f} | "
|
||||
f"LR: {current_lr:.6f} | "
|
||||
f"Time: {epoch_time:.1f}s"
|
||||
)
|
||||
|
||||
# 保存最佳模型
|
||||
if val_score > best_score:
|
||||
best_score = val_score
|
||||
best_path = os.path.join(output_dir, "mask_rcnn_best.pth")
|
||||
save_checkpoint(model, optimizer, scheduler, epoch + 1, label_map,
|
||||
loss_history, best_path)
|
||||
print(f" -> 新最佳模型已保存 (score={best_score:.4f})")
|
||||
|
||||
# 定期保存检查点
|
||||
if (epoch + 1) % args.save_freq == 0 or epoch + 1 == args.epochs:
|
||||
ckpt_path = os.path.join(output_dir, f"mask_rcnn_epoch_{epoch+1}.pth")
|
||||
save_checkpoint(model, optimizer, scheduler, epoch + 1, label_map,
|
||||
loss_history, ckpt_path)
|
||||
print(f" -> 检查点已保存: {ckpt_path}")
|
||||
|
||||
# 保存最终模型
|
||||
final_path = os.path.join(output_dir, "mask_rcnn_final.pth")
|
||||
save_checkpoint(model, optimizer, scheduler, args.epochs, label_map,
|
||||
loss_history, final_path)
|
||||
print(f"\n训练完成！最终模型: {final_path}")
|
||||
|
||||
# 保存损失历史到 JSON
|
||||
with open(os.path.join(output_dir, "loss_history.json"), "w", encoding="utf-8") as f:
|
||||
json.dump(loss_history, f, ensure_ascii=False, indent=2)
|
||||
|
||||
# 绘制损失曲线
|
||||
try:
|
||||
_plot_loss_curve(loss_history, output_dir)
|
||||
except Exception as e:
|
||||
print(f"警告: 无法绘制损失曲线 ({e})")
|
||||
|
||||
print(f"最佳验证分数: {best_score:.4f}")
|
||||
|
||||
|
||||
def _plot_loss_curve(loss_history: List[Dict], output_dir: str):
|
||||
"""绘制并保存损失曲线图"""
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
epochs = [h["epoch"] for h in loss_history]
|
||||
total_loss = [h["train_loss"] for h in loss_history]
|
||||
cls_loss = [h["loss_classifier"] for h in loss_history]
|
||||
box_loss = [h["loss_box_reg"] for h in loss_history]
|
||||
mask_loss = [h["loss_mask"] for h in loss_history]
|
||||
val_scores = [h.get("val_score", 0) for h in loss_history]
|
||||
|
||||
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
|
||||
|
||||
# 损失曲线
|
||||
ax1 = axes[0]
|
||||
ax1.plot(epochs, total_loss, "b-", label="Total Loss", linewidth=2)
|
||||
ax1.plot(epochs, cls_loss, "r--", label="Classifier Loss")
|
||||
ax1.plot(epochs, box_loss, "g--", label="Box Reg Loss")
|
||||
ax1.plot(epochs, mask_loss, "m--", label="Mask Loss")
|
||||
ax1.set_xlabel("Epoch")
|
||||
ax1.set_ylabel("Loss")
|
||||
ax1.set_title("Training Loss")
|
||||
ax1.legend()
|
||||
ax1.grid(True, alpha=0.3)
|
||||
|
||||
# 验证分数曲线
|
||||
ax2 = axes[1]
|
||||
ax2.plot(epochs, val_scores, "b-o", label="Val Score")
|
||||
ax2.set_xlabel("Epoch")
|
||||
ax2.set_ylabel("Score")
|
||||
ax2.set_title("Validation Score")
|
||||
ax2.legend()
|
||||
ax2.grid(True, alpha=0.3)
|
||||
|
||||
plt.tight_layout()
|
||||
curve_path = os.path.join(output_dir, "loss_curve.png")
|
||||
plt.savefig(curve_path, dpi=150, bbox_inches="tight")
|
||||
plt.close()
|
||||
print(f"损失曲线图已保存: {curve_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
|
||||
|
||||
"""
|
||||
# 训练（小数据集推荐冻结骨干网络）
|
||||
python train.py --freeze-backbone --epochs 100
|
||||
|
||||
# 推理
|
||||
python predict.py --image dataset/1.jpg --checkpoint checkpoints/mask_rcnn_best.pth
|
||||
|
||||
# 批量推理 + GT 对比
|
||||
python predict.py --dataset dataset --show-gt
|
||||
|
||||
"""
|
||||
@@ -0,0 +1,464 @@
|
||||
"""
|
||||
工具函数模块
|
||||
包含: LabelMe 标注解析、linestrip 合并为多边形、多边形转 mask、可视化等
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import math
|
||||
import random
|
||||
from typing import List, Dict, Tuple, Optional
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 1. 标签映射
|
||||
# ============================================================
|
||||
|
||||
def get_label_map(dataset_dir: str) -> Dict[str, int]:
|
||||
"""
|
||||
扫描数据集中所有 LabelMe JSON 文件，自动构建标签名→类别ID 映射。
|
||||
类别从 1 开始（0 保留给背景）。
|
||||
|
||||
Args:
|
||||
dataset_dir: 数据集目录路径
|
||||
|
||||
Returns:
|
||||
标签名 → 类别ID 的字典，如 {"crosswalk": 1}
|
||||
"""
|
||||
labels = set()
|
||||
for fname in sorted(os.listdir(dataset_dir)):
|
||||
if not fname.endswith(".json"):
|
||||
continue
|
||||
with open(os.path.join(dataset_dir, fname), "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
for shape in data.get("shapes", []):
|
||||
labels.add(shape["label"])
|
||||
|
||||
# 排序保证可复现
|
||||
label_map = {label: idx + 1 for idx, label in enumerate(sorted(labels))}
|
||||
return label_map
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 2. Linestrip 合并为多边形
|
||||
# ============================================================
|
||||
|
||||
def _point_dist(p1, p2) -> float:
|
||||
"""计算两点之间的欧氏距离"""
|
||||
return math.sqrt((p1[0] - p2[0]) ** 2 + (p1[1] - p2[1]) ** 2)
|
||||
|
||||
|
||||
def merge_linestrips_to_polygon(strips: List[List[List[float]]]) -> List[List[float]]:
|
||||
"""
|
||||
将多条 linestrip 合并为一个闭合多边形。
|
||||
|
||||
策略：以第一条 strip 为起点，每次找到与当前多边形末端点距离最近的
|
||||
下一条 strip 的端点，将该 strip（可能反转）拼接到多边形末尾，
|
||||
最后闭合多边形。
|
||||
|
||||
对于单条 linestrip，直接闭合。
|
||||
|
||||
Args:
|
||||
strips: 多条 linestrip 的点列表，每条 strip 是 [[x, y], ...]
|
||||
|
||||
Returns:
|
||||
闭合多边形的点列表 [[x, y], ..., [x, y]]（首尾相同）
|
||||
"""
|
||||
if len(strips) == 0:
|
||||
return []
|
||||
|
||||
# 深拷贝避免修改原始数据
|
||||
strips = [list(s) for s in strips]
|
||||
|
||||
if len(strips) == 1:
|
||||
polygon = list(strips[0])
|
||||
polygon.append(list(polygon[0])) # 闭合
|
||||
return polygon
|
||||
|
||||
# 从第一条 strip 开始
|
||||
polygon = list(strips[0])
|
||||
remaining = strips[1:]
|
||||
|
||||
while remaining:
|
||||
last_point = polygon[-1]
|
||||
best_idx = 0
|
||||
best_dist = float("inf")
|
||||
best_reverse = False
|
||||
|
||||
for i, strip in enumerate(remaining):
|
||||
# 尝试正向连接（strip 的起点与 polygon 末端最近）
|
||||
d_start = _point_dist(last_point, strip[0])
|
||||
# 尝试反向连接（strip 的终点与 polygon 末端最近）
|
||||
d_end = _point_dist(last_point, strip[-1])
|
||||
|
||||
if d_start < best_dist:
|
||||
best_dist = d_start
|
||||
best_idx = i
|
||||
best_reverse = False
|
||||
if d_end < best_dist:
|
||||
best_dist = d_end
|
||||
best_idx = i
|
||||
best_reverse = True
|
||||
|
||||
strip = remaining.pop(best_idx)
|
||||
if best_reverse:
|
||||
polygon.extend(list(reversed(strip)))
|
||||
else:
|
||||
polygon.extend(strip)
|
||||
|
||||
# 闭合多边形
|
||||
polygon.append(list(polygon[0]))
|
||||
return polygon
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 3. 多边形 / Linestrip → Mask
|
||||
# ============================================================
|
||||
|
||||
def shapes_to_masks(
|
||||
shapes: List[dict],
|
||||
height: int,
|
||||
width: int,
|
||||
label_map: Dict[str, int],
|
||||
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
"""
|
||||
将 LabelMe 的 shapes 转换为实例分割所需的 masks、labels、boxes。
|
||||
|
||||
处理逻辑：
|
||||
- polygon 类型：直接作为多边形填充
|
||||
- linestrip 类型：将相同 label 的 linestrip 合并为一个多边形后填充
|
||||
- rectangle / circle 类型：转换为多边形后处理
|
||||
|
||||
Args:
|
||||
shapes: LabelMe JSON 中的 shapes 列表
|
||||
height: 图像高度
|
||||
width: 图像宽度
|
||||
label_map: 标签名→类别ID 映射
|
||||
|
||||
Returns:
|
||||
masks: [N, H, W] uint8 二值 mask（0 或 1）
|
||||
labels: [N] int64 类别ID
|
||||
boxes: [N, 4] float32 格式 (x1, y1, x2, y2)
|
||||
"""
|
||||
# 按 label 分组 linestrip（同 label 的 linestrip 合并为一个实例）
|
||||
linestrip_groups: Dict[str, List[List]] = {}
|
||||
# 非 linestrip 的 shape 逐个处理
|
||||
other_shapes: List[dict] = []
|
||||
|
||||
for shape in shapes:
|
||||
shape_type = shape.get("shape_type", "polygon")
|
||||
if shape_type == "linestrip":
|
||||
label = shape["label"]
|
||||
linestrip_groups.setdefault(label, []).append(shape["points"])
|
||||
else:
|
||||
other_shapes.append(shape)
|
||||
|
||||
all_masks = []
|
||||
all_labels = []
|
||||
|
||||
# 处理合并后的 linestrip
|
||||
for label, strips in linestrip_groups.items():
|
||||
polygon = merge_linestrips_to_polygon(strips)
|
||||
if len(polygon) < 3:
|
||||
continue
|
||||
mask = _polygon_to_mask(polygon, height, width)
|
||||
if mask.sum() == 0:
|
||||
continue
|
||||
all_masks.append(mask)
|
||||
all_labels.append(label_map.get(label, 1))
|
||||
|
||||
# 处理其他 shape 类型
|
||||
for shape in other_shapes:
|
||||
label = shape["label"]
|
||||
shape_type = shape.get("shape_type", "polygon")
|
||||
points = shape["points"]
|
||||
|
||||
if shape_type == "polygon":
|
||||
polygon = points
|
||||
elif shape_type == "rectangle":
|
||||
# 矩形转多边形
|
||||
x1, y1 = points[0]
|
||||
x2, y2 = points[1]
|
||||
polygon = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]
|
||||
elif shape_type == "circle":
|
||||
# 圆形转多边形（近似）
|
||||
cx, cy = points[0]
|
||||
rx = points[1][0] - cx
|
||||
ry = points[1][1] - cy
|
||||
polygon = []
|
||||
for angle in range(0, 360, 5):
|
||||
rad = math.radians(angle)
|
||||
polygon.append([cx + rx * math.cos(rad), cy + ry * math.sin(rad)])
|
||||
else:
|
||||
continue
|
||||
|
||||
if len(polygon) < 3:
|
||||
continue
|
||||
|
||||
mask = _polygon_to_mask(polygon, height, width)
|
||||
if mask.sum() == 0:
|
||||
continue
|
||||
all_masks.append(mask)
|
||||
all_labels.append(label_map.get(label, 1))
|
||||
|
||||
if len(all_masks) == 0:
|
||||
return (
|
||||
np.zeros((0, height, width), dtype=np.uint8),
|
||||
np.zeros((0,), dtype=np.int64),
|
||||
np.zeros((0, 4), dtype=np.float32),
|
||||
)
|
||||
|
||||
masks = np.stack(all_masks, axis=0) # [N, H, W]
|
||||
labels = np.array(all_labels, dtype=np.int64)
|
||||
|
||||
# 从 mask 计算包围框
|
||||
boxes = masks_to_boxes(masks)
|
||||
|
||||
return masks, labels, boxes
|
||||
|
||||
|
||||
def _polygon_to_mask(
|
||||
polygon: List[List[float]],
|
||||
height: int,
|
||||
width: int,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
将多边形转换为二值 mask。
|
||||
|
||||
Args:
|
||||
polygon: 多边形点列表 [[x, y], ...]
|
||||
height: 图像高度
|
||||
width: 图像宽度
|
||||
|
||||
Returns:
|
||||
[H, W] uint8 二值 mask（0 或 1）
|
||||
"""
|
||||
img = Image.new("L", (width, height), 0)
|
||||
ImageDraw.Draw(img).polygon(
|
||||
[(float(p[0]), float(p[1])) for p in polygon],
|
||||
fill=1,
|
||||
outline=1,
|
||||
)
|
||||
return np.array(img, dtype=np.uint8)
|
||||
|
||||
|
||||
def masks_to_boxes(masks: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
从一组二值 mask 计算每个实例的包围框。
|
||||
|
||||
Args:
|
||||
masks: [N, H, W] uint8
|
||||
|
||||
Returns:
|
||||
[N, 4] float32, 格式 (x1, y1, x2, y2)
|
||||
"""
|
||||
boxes = np.zeros((masks.shape[0], 4), dtype=np.float32)
|
||||
for i, mask in enumerate(masks):
|
||||
ys, xs = np.where(mask > 0)
|
||||
if len(xs) == 0:
|
||||
continue
|
||||
boxes[i] = [xs.min(), ys.min(), xs.max(), ys.max()]
|
||||
return boxes
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 4. 可视化
|
||||
# ============================================================
|
||||
|
||||
def visualize_instances(
|
||||
image: np.ndarray,
|
||||
boxes: np.ndarray,
|
||||
masks: np.ndarray,
|
||||
labels: np.ndarray,
|
||||
scores: Optional[np.ndarray] = None,
|
||||
label_map: Optional[Dict[int, str]] = None,
|
||||
score_threshold: float = 0.5,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
在图像上绘制实例分割结果（包围框 + mask 叠加）。
|
||||
|
||||
Args:
|
||||
image: [H, W, 3] uint8 BGR 或 RGB 图像
|
||||
boxes: [N, 4] 包围框
|
||||
masks: [N, H, W] 二值 mask
|
||||
labels: [N] 类别ID
|
||||
scores: [N] 置信度（可选）
|
||||
label_map: 类别ID→名称映射（可选）
|
||||
score_threshold: 置信度阈值
|
||||
|
||||
Returns:
|
||||
绘制后的图像 [H, W, 3] uint8
|
||||
"""
|
||||
try:
|
||||
import cv2
|
||||
except ImportError:
|
||||
# 如果没有 cv2，退化为 PIL 绘制
|
||||
return _visualize_instances_pil(
|
||||
image, boxes, masks, labels, scores, label_map, score_threshold
|
||||
)
|
||||
|
||||
img = image.copy()
|
||||
if img.dtype != np.uint8:
|
||||
img = (img * 255).astype(np.uint8)
|
||||
|
||||
# 生成固定颜色
|
||||
np.random.seed(42)
|
||||
num_classes = max(int(labels.max()) + 1, 10) if len(labels) > 0 else 10
|
||||
colors = np.random.randint(0, 255, size=(num_classes, 3), dtype=np.uint8)
|
||||
|
||||
for i in range(len(boxes)):
|
||||
if scores is not None and scores[i] < score_threshold:
|
||||
continue
|
||||
|
||||
color = colors[int(labels[i])].tolist()
|
||||
x1, y1, x2, y2 = boxes[i].astype(int)
|
||||
|
||||
# 绘制 mask 半透明叠加
|
||||
mask = masks[i]
|
||||
colored_mask = np.zeros_like(img)
|
||||
for c in range(3):
|
||||
colored_mask[:, :, c] = mask * color[c]
|
||||
img = cv2.addWeighted(img, 1.0, colored_mask, 0.5, 0)
|
||||
|
||||
# 绘制包围框
|
||||
cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
|
||||
|
||||
# 绘制标签
|
||||
label_text = label_map.get(int(labels[i]), str(labels[i])) if label_map else str(labels[i])
|
||||
if scores is not None:
|
||||
label_text = f"{label_text}: {scores[i]:.2f}"
|
||||
|
||||
(tw, th), _ = cv2.getTextSize(label_text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1)
|
||||
cv2.rectangle(img, (x1, y1 - th - 6), (x1 + tw + 4, y1), color, -1)
|
||||
cv2.putText(
|
||||
img, label_text, (x1 + 2, y1 - 4),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 1, cv2.LINE_AA,
|
||||
)
|
||||
|
||||
return img
|
||||
|
||||
|
||||
def _visualize_instances_pil(
|
||||
image: np.ndarray,
|
||||
boxes: np.ndarray,
|
||||
masks: np.ndarray,
|
||||
labels: np.ndarray,
|
||||
scores: Optional[np.ndarray],
|
||||
label_map: Optional[Dict[int, str]],
|
||||
score_threshold: float,
|
||||
) -> np.ndarray:
|
||||
"""使用 PIL 的可视化后备方案（当 cv2 不可用时）"""
|
||||
from PIL import Image as PILImage, ImageDraw as PILDraw, ImageFont
|
||||
|
||||
if image.dtype != np.uint8:
|
||||
image = (image * 255).astype(np.uint8)
|
||||
|
||||
# 转 RGB
|
||||
if image.shape[2] == 3:
|
||||
pil_img = PILImage.fromarray(image[..., ::-1]) if image.dtype == np.uint8 else PILImage.fromarray(image)
|
||||
else:
|
||||
pil_img = PILImage.fromarray(image)
|
||||
|
||||
np.random.seed(42)
|
||||
num_classes = max(int(labels.max()) + 1, 10) if len(labels) > 0 else 10
|
||||
colors = np.random.randint(0, 255, size=(num_classes, 3), dtype=np.uint8)
|
||||
|
||||
overlay = pil_img.copy()
|
||||
draw_overlay = PILDraw.Draw(overlay)
|
||||
draw = PILDraw.Draw(pil_img)
|
||||
|
||||
for i in range(len(boxes)):
|
||||
if scores is not None and scores[i] < score_threshold:
|
||||
continue
|
||||
|
||||
color = tuple(colors[int(labels[i])].tolist())
|
||||
mask = masks[i]
|
||||
# 用颜色填充 mask 区域
|
||||
mask_img = PILImage.new("RGB", pil_img.size, (0, 0, 0))
|
||||
mask_draw = PILDraw.Draw(mask_img)
|
||||
mask_pil = PILImage.fromarray((mask * 255).astype(np.uint8))
|
||||
# 叠加半透明 mask
|
||||
for c in range(3):
|
||||
channel = np.array(pil_img)[:, :, c].astype(float)
|
||||
channel[mask > 0] = channel[mask > 0] * 0.5 + color[c] * 0.5
|
||||
pil_img_array = np.array(pil_img)
|
||||
pil_img_array[:, :, c] = channel
|
||||
pil_img = PILImage.fromarray(pil_img_array.astype(np.uint8))
|
||||
|
||||
draw = PILDraw.Draw(pil_img)
|
||||
for i in range(len(boxes)):
|
||||
if scores is not None and scores[i] < score_threshold:
|
||||
continue
|
||||
color = tuple(colors[int(labels[i])].tolist())
|
||||
x1, y1, x2, y2 = boxes[i].astype(int)
|
||||
draw.rectangle([x1, y1, x2, y2], outline=color, width=2)
|
||||
|
||||
label_text = label_map.get(int(labels[i]), str(labels[i])) if label_map else str(labels[i])
|
||||
if scores is not None:
|
||||
label_text = f"{label_text}: {scores[i]:.2f}"
|
||||
draw.text((x1 + 2, max(0, y1 - 15)), label_text, fill=color)
|
||||
|
||||
return np.array(pil_img)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 5. 数据集划分
|
||||
# ============================================================
|
||||
|
||||
def get_image_json_pairs(dataset_dir: str) -> List[Tuple[str, str]]:
|
||||
"""
|
||||
扫描数据集目录，返回 (image_path, json_path) 配对列表。
|
||||
|
||||
Args:
|
||||
dataset_dir: 数据集目录
|
||||
|
||||
Returns:
|
||||
[(image_path, json_path), ...]
|
||||
"""
|
||||
pairs = []
|
||||
for fname in sorted(os.listdir(dataset_dir)):
|
||||
if not fname.endswith(".json"):
|
||||
continue
|
||||
json_path = os.path.join(dataset_dir, fname)
|
||||
with open(json_path, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
image_name = data.get("imagePath", fname.replace(".json", ".jpg"))
|
||||
image_path = os.path.join(dataset_dir, image_name)
|
||||
if not os.path.exists(image_path):
|
||||
# 尝试其他常见扩展名
|
||||
for ext in [".jpg", ".jpeg", ".png", ".bmp"]:
|
||||
alt_path = os.path.join(dataset_dir, fname.replace(".json", ext))
|
||||
if os.path.exists(alt_path):
|
||||
image_path = alt_path
|
||||
break
|
||||
pairs.append((image_path, json_path))
|
||||
return pairs
|
||||
|
||||
|
||||
def split_dataset(
|
||||
pairs: List[Tuple[str, str]],
|
||||
val_ratio: float = 0.2,
|
||||
seed: int = 42,
|
||||
) -> Tuple[List, List]:
|
||||
"""
|
||||
将数据集划分为训练集和验证集。
|
||||
|
||||
Args:
|
||||
pairs: (image_path, json_path) 配对列表
|
||||
val_ratio: 验证集比例
|
||||
seed: 随机种子
|
||||
|
||||
Returns:
|
||||
train_pairs, val_pairs
|
||||
"""
|
||||
rng = random.Random(seed)
|
||||
indices = list(range(len(pairs)))
|
||||
rng.shuffle(indices)
|
||||
val_size = max(1, int(len(pairs) * val_ratio)) if len(pairs) > 1 else 0
|
||||
val_indices = set(indices[:val_size])
|
||||
train_pairs = [pairs[i] for i in indices if i not in val_indices]
|
||||
val_pairs = [pairs[i] for i in val_indices]
|
||||
return train_pairs, val_pairs
|
||||
Reference in New Issue
Block a user