175 lines
5.1 KiB
Markdown
175 lines
5.1 KiB
Markdown
# YOLO 分割模型训练说明
|
||
|
||
## 概述
|
||
|
||
本项目现在支持使用多边形标注进行实例分割训练，可以精确框选目标，避免矩形框包含非目标区域的问题。
|
||
|
||
## 文件说明
|
||
|
||
### 1. `train_segmentation.py` - 分割模型训练脚本
|
||
专门用于训练YOLO分割模型（支持多边形标注）
|
||
|
||
**使用方法：**
|
||
```bash
|
||
python train_segmentation.py
|
||
```
|
||
|
||
**参数说明：**
|
||
- `dataset_path`: 数据集配置文件路径（默认: `datasets/搭电设备/data.yaml`）
|
||
- `model_name`: 预训练模型（默认: `yolov8n-seg.pt`）
|
||
- 可选: `yolov8n-seg.pt`, `yolov8s-seg.pt`, `yolov8m-seg.pt`, `yolov8l-seg.pt`, `yolov8x-seg.pt`
|
||
- 或: `yolo11n-seg.pt`, `yolo11s-seg.pt` 等
|
||
- `epochs`: 训练轮数（默认: 100）
|
||
- `batch`: 批次大小（默认: 8）
|
||
- `imgsz`: 图片尺寸（默认: 640）
|
||
- `device`: 设备类型（默认: `"cpu"`，GPU使用 `"cuda"` 或 `"0"`）
|
||
|
||
### 2. `verify_seg_labels.py` - 标注格式验证工具
|
||
批量验证YOLO分割标注格式是否正确
|
||
|
||
**使用方法：**
|
||
```bash
|
||
# 验证指定目录的标注文件
|
||
python verify_seg_labels.py datasets/搭电设备/train/labels datasets/搭电设备/train/images
|
||
|
||
# 或直接运行（会验证搭电设备数据集）
|
||
python verify_seg_labels.py
|
||
```
|
||
|
||
**功能：**
|
||
- 检查标注格式是否正确
|
||
- 验证坐标是否在有效范围内 [0, 1]
|
||
- 统计多边形顶点数量
|
||
- 可选：可视化标注结果
|
||
|
||
### 3. `test_seg_label.py` - 单个文件测试工具
|
||
快速测试单个标注文件
|
||
|
||
**使用方法：**
|
||
```bash
|
||
# 仅验证标注文件
|
||
python test_seg_label.py datasets/搭电设备/train/labels/your_file.txt
|
||
|
||
# 验证并可视化
|
||
python test_seg_label.py datasets/搭电设备/train/labels/your_file.txt datasets/搭电设备/train/images/your_file.jpg
|
||
```
|
||
|
||
### 4. `train_jumper_cable.py` - 智能训练脚本（已更新）
|
||
自动检测标注格式，选择相应的模型
|
||
|
||
- 如果检测到多边形格式（坐标数 > 5），使用分割模型
|
||
- 如果检测到边界框格式（坐标数 = 5），使用检测模型
|
||
|
||
## 标注格式说明
|
||
|
||
### YOLO 分割格式（多边形）
|
||
|
||
每个 `.txt` 文件格式：
|
||
```
|
||
class_id x1 y1 x2 y2 x3 y3 x4 y4 ...
|
||
```
|
||
|
||
**示例：**
|
||
```
|
||
0 0.666693 0.597539 0.600534 0.649248 0.427135 0.689463 0.287852 0.680586 ...
|
||
```
|
||
|
||
**说明：**
|
||
- `class_id`: 类别ID（整数，从0开始）
|
||
- `x1 y1 x2 y2 ...`: 归一化的多边形顶点坐标（0-1之间）
|
||
- 至少需要3个点（6个坐标值）
|
||
|
||
### 从 CVAT 导出
|
||
|
||
1. 在 CVAT 中选择任务/作业
|
||
2. 点击 "Export job as a dataset"
|
||
3. **选择 "Ultralytics YOLO Segmentation 1.0"**
|
||
4. 下载导出的数据集
|
||
|
||
## 数据集目录结构
|
||
|
||
```
|
||
datasets/
|
||
└── 搭电设备/
|
||
├── data.yaml # 数据集配置
|
||
├── train/
|
||
│ ├── images/ # 训练图片
|
||
│ └── labels/ # 训练标注（.txt文件）
|
||
└── val/
|
||
├── images/ # 验证图片
|
||
└── labels/ # 验证标注（.txt文件）
|
||
```
|
||
|
||
## data.yaml 配置示例
|
||
|
||
```yaml
|
||
train: train/images
|
||
val: val/images
|
||
nc: 2
|
||
names: ['JumperCable', 'NoJumperCable']
|
||
```
|
||
|
||
## 训练步骤
|
||
|
||
### 1. 准备数据
|
||
- 从 CVAT 导出为 "Ultralytics YOLO Segmentation 1.0" 格式
|
||
- 将数据组织到正确的目录结构
|
||
- 创建或更新 `data.yaml` 配置文件
|
||
|
||
### 2. 验证标注
|
||
```bash
|
||
# 验证标注格式
|
||
python verify_seg_labels.py datasets/搭电设备/train/labels datasets/搭电设备/train/images
|
||
```
|
||
|
||
### 3. 开始训练
|
||
```bash
|
||
# 使用分割模型训练
|
||
python train_segmentation.py
|
||
```
|
||
|
||
或修改 `train_segmentation.py` 中的参数：
|
||
```python
|
||
train_segmentation_model(
|
||
dataset_path="datasets/你的数据集/data.yaml",
|
||
model_name="yolov8n-seg.pt",
|
||
epochs=100,
|
||
batch=8,
|
||
imgsz=640,
|
||
device="cuda" # 如果有GPU
|
||
)
|
||
```
|
||
|
||
### 4. 查看结果
|
||
训练完成后，模型保存在 `runs/segment/train/weights/` 目录下：
|
||
- `best.pt`: 最佳模型
|
||
- `last.pt`: 最后一轮的模型
|
||
|
||
## 注意事项
|
||
|
||
1. **模型选择**：必须使用分割模型（`-seg` 后缀），如 `yolov8n-seg.pt`
|
||
2. **标注格式**：确保标注是多边形格式（坐标数 > 5）
|
||
3. **文件名对应**：标注文件名必须与图片文件名一致（仅扩展名不同）
|
||
4. **坐标归一化**：所有坐标必须在 [0, 1] 范围内
|
||
5. **GPU加速**：如果有GPU，将 `device` 设置为 `"cuda"` 或 `"0"` 可以大幅加速训练
|
||
|
||
## 常见问题
|
||
|
||
### Q: 如何判断标注是分割格式还是检测格式？
|
||
A: 分割格式的坐标数量 > 5（class_id + 至少6个坐标值），检测格式只有5个值（class_id + 4个边界框坐标）
|
||
|
||
### Q: 训练时提示找不到模型文件？
|
||
A: 首次使用时会自动下载，确保网络连接正常。也可以手动下载后放到项目根目录。
|
||
|
||
### Q: 如何可视化分割结果？
|
||
A: 使用 `test_seg_label.py` 脚本，提供图片路径即可可视化标注。
|
||
|
||
### Q: 训练速度很慢？
|
||
A: 如果有GPU，将 `device="cpu"` 改为 `device="cuda"`。也可以减小 `batch` 大小或 `imgsz` 尺寸。
|
||
|
||
## 参考
|
||
|
||
- [Ultralytics YOLO 文档](https://docs.ultralytics.com/)
|
||
- [YOLO 分割任务说明](https://docs.ultralytics.com/tasks/segment/)
|
||
|