# 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工程索引

相关笔记：
- [[YOLO+ByteTrack水流速度检测-可行性详细分析]] — **完整可行性分析与实现指南（2026-04-23 更新）**
- [[INDEX_AI工程]] - AI工程知识索引
- [[ClaudeCode完全研究]] - Claude Code 完整指南

---

*整理：知识库管理员 | 归档：2026-04-23*