commit 61f6ea918e798ac76cf5d6dca2fe0d298f2bebcb
Author: flk <2640748194@qq.com>
Date:   Mon Aug 10 01:20:28 2026 +0800

    Mask R-CNN instance segmentation training code

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..9e259e4
--- /dev/null
+++ b/.gitignore
@@ -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
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..f1b2726
--- /dev/null
+++ b/README.md
@@ -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）
diff --git a/dataset.py b/dataset.py
new file mode 100644
index 0000000..2f39929
--- /dev/null
+++ b/dataset.py
@@ -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([])
diff --git a/dataset/1.jpg b/dataset/1.jpg
new file mode 100644
index 0000000..337acc8
Binary files /dev/null and b/dataset/1.jpg differ
diff --git a/dataset/1.json b/dataset/1.json
new file mode 100644
index 0000000..16f52b7
--- /dev/null
+++ b/dataset/1.json
@@ -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
+}
\ No newline at end of file
diff --git a/dataset/2.jpg b/dataset/2.jpg
new file mode 100644
index 0000000..dad3fa3
Binary files /dev/null and b/dataset/2.jpg differ
diff --git a/dataset/2.json b/dataset/2.json
new file mode 100644
index 0000000..1e02fb1
--- /dev/null
+++ b/dataset/2.json
@@ -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
+}
\ No newline at end of file
diff --git a/dataset/3.jpg b/dataset/3.jpg
new file mode 100644
index 0000000..86cb07f
Binary files /dev/null and b/dataset/3.jpg differ
diff --git a/dataset/3.json b/dataset/3.json
new file mode 100644
index 0000000..d879b01
--- /dev/null
+++ b/dataset/3.json
@@ -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
+}
\ No newline at end of file
diff --git a/dataset/4.jpg b/dataset/4.jpg
new file mode 100644
index 0000000..5effc62
Binary files /dev/null and b/dataset/4.jpg differ
diff --git a/dataset/4.json b/dataset/4.json
new file mode 100644
index 0000000..f0568ee
--- /dev/null
+++ b/dataset/4.json
@@ -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
+}
\ No newline at end of file
diff --git a/model.py b/model.py
new file mode 100644
index 0000000..3e020d5
--- /dev/null
+++ b/model.py
@@ -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)
diff --git a/predict.py b/predict.py
new file mode 100644
index 0000000..29d7123
--- /dev/null
+++ b/predict.py
@@ -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()
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..a5a2582
--- /dev/null
+++ b/requirements.txt
@@ -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
diff --git a/train.py b/train.py
new file mode 100644
index 0000000..0fea04a
--- /dev/null
+++ b/train.py
@@ -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
+
+"""
diff --git a/utils.py b/utils.py
new file mode 100644
index 0000000..639eddc
--- /dev/null
+++ b/utils.py
@@ -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