diff --git a/AI工程/YOLO+ByteTrack水流速度检测.md b/AI工程/YOLO+ByteTrack水流速度检测.md new file mode 100755 index 0000000..12edca0 --- /dev/null +++ b/AI工程/YOLO+ByteTrack水流速度检测.md @@ -0,0 +1,310 @@ +# 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* \ No newline at end of file