310 lines
7.0 KiB
Markdown
Executable File
310 lines
7.0 KiB
Markdown
Executable File
# YOLO + ByteTrack 水流速度检测模型
|
||
|
||
> 基于 YOLO 目标检测 + ByteTrack 多目标追踪的水流速度检测方案
|
||
>
|
||
> 归档：2026-04-23
|
||
|
||
---
|
||
|
||
## 🎯 项目概述
|
||
|
||
**核心思路**：
|
||
- 用 YOLO 检测水流中的物体（如漂浮物、示踪粒子等）
|
||
- 用 ByteTrack 追踪每个物体的运动轨迹
|
||
- 通过轨迹位移和时间计算流速
|
||
|
||
---
|
||
|
||
## 🏗️ 技术架构
|
||
|
||
```
|
||
视频输入 → YOLO 检测 → 边界框 + 置信度
|
||
↓
|
||
ByteTrack 追踪关联
|
||
↓
|
||
物体 ID + 轨迹
|
||
↓
|
||
位移 / 时间 = 流速
|
||
```
|
||
|
||
---
|
||
|
||
## 📦 核心组件
|
||
|
||
### 1. YOLO 检测器
|
||
|
||
| 选项 | 说明 |
|
||
|------|------|
|
||
| **YOLOX** | ByteTrack 原生检测器，anchor-free，推荐 |
|
||
| **YOLOv7** | 性能更强，训练更灵活 |
|
||
| **YOLOv8** | 最新 Ultralytics 版本，生态最好 |
|
||
| **YOLOv5** | 最流行，社区支持广 |
|
||
|
||
**推荐**：YOLOX 或 YOLOv8（与 ByteTrack 集成好）
|
||
|
||
### 2. ByteTrack 追踪器
|
||
|
||
| 项目 | 数据 |
|
||
|------|------|
|
||
| 论文 | ECCV 2022 |
|
||
| GitHub | ifzhang/ByteTrack |
|
||
| 原理 | 关联所有检测框（不只是高分框） |
|
||
| 速度 | 30 FPS (V100) |
|
||
| 性能 | MOT17: 80.3 MOTA, 77.3 IDF1 |
|
||
|
||
---
|
||
|
||
## 🔧 训练自定义模型
|
||
|
||
### Step 1: 数据准备
|
||
|
||
**需要**：
|
||
- 水流场景视频/图像
|
||
- 标注目标物体（如漂浮物、示踪粒子）
|
||
- COCO 格式标注
|
||
|
||
**标注工具**：
|
||
- LabelImg
|
||
- CVAT
|
||
- Roboflow
|
||
- Label Studio
|
||
|
||
**COCO 格式**：
|
||
```json
|
||
{
|
||
"images": [{"id": 1, "file_name": "frame_001.jpg", "width": 1920, "height": 1080}],
|
||
"annotations": [{"id": 1, "image_id": 1, "category_id": 1, "bbox": [x, y, w, h]}],
|
||
"categories": [{"id": 1, "name": "floating_object"}]
|
||
}
|
||
```
|
||
|
||
### Step 2: YOLO 训练
|
||
|
||
**以 YOLOX 为例**：
|
||
```bash
|
||
# 1. 克隆仓库
|
||
git clone https://github.com/Megvii-BaseDetection/YOLOX.git
|
||
cd YOLOX
|
||
pip3 install -v -e .
|
||
|
||
# 2. 准备数据集（COCO 格式）
|
||
ln -s /path/to/your/data ./datasets/water_flow
|
||
|
||
# 3. 创建配置文件
|
||
# 复制 exps/example/yolox_voc/yolox_voc_s.py 并修改
|
||
|
||
# 4. 训练
|
||
python tools/train.py -f exps/example/water_flow/yolox_s_water.py \
|
||
-d 1 -b 8 --fp16 -o -c pretrained/yolox_s.pth
|
||
```
|
||
|
||
**以 YOLOv7 为例**：
|
||
```bash
|
||
# 1. 克隆仓库
|
||
git clone https://github.com/WongKinYiu/yolov7.git
|
||
cd yolov7
|
||
|
||
# 2. 准备数据（YOLO 格式）
|
||
# 创建 data/custom.yaml
|
||
|
||
# 3. 微调训练
|
||
python train.py --device 0 --batch-size 32 \
|
||
--data data/custom.yaml --img 640 640 \
|
||
--cfg cfg/training/yolov7-custom.yaml \
|
||
--weights 'yolov7_training.pt' --name custom
|
||
```
|
||
|
||
### Step 3: ByteTrack 集成
|
||
|
||
```bash
|
||
# 1. 克隆 ByteTrack
|
||
git clone https://github.com/ifzhang/ByteTrack.git
|
||
cd ByteTrack
|
||
pip3 install -r requirements.txt
|
||
python3 setup.py develop
|
||
|
||
# 2. 训练 ByteTrack 模型（使用自定义数据）
|
||
python3 tools/train.py -f exps/example/mot/your_exp_file.py \
|
||
-d 8 -b 48 --fp16 -o -c pretrained/yolox_x.pth
|
||
|
||
# 3. 追踪推理
|
||
python3 tools/track.py -f exps/example/mot/your_exp_file.py \
|
||
-c pretrained/bytetrack_custom.pth.tar -b 1 -d 1 --fp16 --fuse
|
||
```
|
||
|
||
---
|
||
|
||
## 📐 流速计算算法
|
||
|
||
### 基本原理
|
||
|
||
```
|
||
v = d / t
|
||
```
|
||
|
||
- `v` = 流速 (m/s)
|
||
- `d` = 实际位移 (m)
|
||
- `t` = 时间间隔 (s)
|
||
|
||
### 像素到实际距离转换
|
||
|
||
**需要相机标定**：
|
||
1. 在场景中放置已知尺寸的参考物
|
||
2. 计算像素到实际距离的比例因子
|
||
3. 考虑透视变形（如果需要）
|
||
|
||
### 实现代码
|
||
|
||
```python
|
||
import cv2
|
||
import numpy as np
|
||
from collections import defaultdict
|
||
|
||
class FlowVelocityCalculator:
|
||
def __init__(self, pixels_per_meter, fps):
|
||
self.ppm = pixels_per_meter # 像素/米 比例
|
||
self.fps = fps # 帧率
|
||
self.trajectories = defaultdict(list)
|
||
|
||
def update(self, tracks, frame_idx):
|
||
"""更新轨迹"""
|
||
for track in tracks:
|
||
track_id = track.track_id
|
||
bbox = track.tlbr # [x1, y1, x2, y2]
|
||
center = ((bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2)
|
||
self.trajectories[track_id].append((frame_idx, center))
|
||
|
||
def calculate_velocity(self, track_id, min_frames=5):
|
||
"""计算单个物体的流速"""
|
||
traj = self.trajectories.get(track_id, [])
|
||
if len(traj) < min_frames:
|
||
return None
|
||
|
||
# 取首尾帧
|
||
start_frame, start_pos = traj[0]
|
||
end_frame, end_pos = traj[-1]
|
||
|
||
# 计算像素位移
|
||
dx = end_pos[0] - start_pos[0]
|
||
dy = end_pos[1] - start_pos[1]
|
||
pixel_dist = np.sqrt(dx**2 + dy**2)
|
||
|
||
# 转换为实际距离
|
||
real_dist = pixel_dist / self.ppm
|
||
|
||
# 计算时间
|
||
time_sec = (end_frame - start_frame) / self.fps
|
||
|
||
# 流速
|
||
velocity = real_dist / time_sec
|
||
return velocity
|
||
|
||
def get_all_velocities(self):
|
||
"""获取所有物体的流速"""
|
||
velocities = {}
|
||
for track_id in self.trajectories:
|
||
v = self.calculate_velocity(track_id)
|
||
if v is not None:
|
||
velocities[track_id] = v
|
||
return velocities
|
||
```
|
||
|
||
---
|
||
|
||
## 📊 完整工作流程
|
||
|
||
```
|
||
1. 数据采集
|
||
↓
|
||
2. 数据标注（COCO 格式）
|
||
↓
|
||
3. YOLO 训练检测器
|
||
↓
|
||
4. ByteTrack 训练追踪器（可选）
|
||
↓
|
||
5. 相机标定（像素→实际距离）
|
||
↓
|
||
6. 视频推理（检测 + 追踪）
|
||
↓
|
||
7. 轨迹提取与流速计算
|
||
↓
|
||
8. 结果可视化与输出
|
||
```
|
||
|
||
---
|
||
|
||
## 🛠️ 环境配置
|
||
|
||
```bash
|
||
# Python 3.8+
|
||
conda create -n waterflow python=3.8
|
||
conda activate waterflow
|
||
|
||
# 安装 PyTorch
|
||
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118
|
||
|
||
# 安装 YOLOX
|
||
git clone https://github.com/Megvii-BaseDetection/YOLOX.git
|
||
cd YOLOX && pip install -v -e .
|
||
|
||
# 安装 ByteTrack
|
||
cd .. && git clone https://github.com/ifzhang/ByteTrack.git
|
||
cd ByteTrack && pip install -r requirements.txt && python setup.py develop
|
||
|
||
# 其他依赖
|
||
pip install opencv-python numpy scipy lap motmetrics
|
||
```
|
||
|
||
---
|
||
|
||
## 📈 模型选择建议
|
||
|
||
| 场景 | 推荐模型 | 原因 |
|
||
|------|---------|------|
|
||
| 边缘部署 | YOLOX-Nano + ByteTrack | 参数 0.9M，速度快 |
|
||
| 平衡性能 | YOLOX-S + ByteTrack | 精度速度均衡 |
|
||
| 高精度 | YOLOX-X + ByteTrack | 精度最高 |
|
||
| 快速验证 | YOLOv8n + ByteTrack | 生态好，易用 |
|
||
|
||
---
|
||
|
||
## 💡 优化技巧
|
||
|
||
### 1. 数据增强
|
||
- 运动模糊模拟
|
||
- 光照变化
|
||
- 水面反射
|
||
|
||
### 2. 追踪优化
|
||
- 调整检测阈值（低阈值可能更好）
|
||
- 调整追踪关联阈值
|
||
- 使用卡尔曼滤波
|
||
|
||
### 3. 流速计算
|
||
- 多物体平均流速
|
||
- 异常值过滤
|
||
- 轨迹平滑处理
|
||
|
||
---
|
||
|
||
## 🔗 相关资源
|
||
|
||
| 资源 | 链接 |
|
||
|------|------|
|
||
| ByteTrack | https://github.com/ifzhang/ByteTrack |
|
||
| YOLOX | https://github.com/Megvii-BaseDetection/YOLOX |
|
||
| YOLOv7 | https://github.com/WongKinYiu/yolov7 |
|
||
| YOLOv8 | https://github.com/ultralytics/ultralytics |
|
||
| ByteTrack 论文 | https://arxiv.org/abs/2110.06864 |
|
||
|
||
---
|
||
|
||
## AI工程索引
|
||
|
||
相关笔记：
|
||
- [[INDEX_AI工程]] - AI工程知识索引
|
||
- [[ClaudeCode完全研究]] - Claude Code 完整指南
|
||
|
||
---
|
||
|
||
*整理：知识库管理员 | 归档：2026-04-23* |