Skip to content

Commit e90d017

Browse files
SEMHAQclaudecharon191910
authored
添加上下行车辆计数版本 (#5229)
* update readme&requirements * feat: 首次提交 - YOLO11n 车辆计数系统 - 创建 main.py 作为项目入口文件,支持命令行参数配置 - 重构 yolo_vehicle_counter.py,支持函数调用和外部参数传入 - 将所有代码注释改为中文并优化说明 - 将路径配置提取到文件开头常量区域,便于管理 - 更新 README.md,中文化并完善项目文档 - 添加中文使用说明和配置说明 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * feat: 重构YOLO11n车辆计数系统 - 重命名目录为yolo11n_vehicle_counter(符合Python命名规范) - 更新README.md,添加模型下载和Google Drive链接说明 - 添加目录创建步骤和.gitignore配置说明 - 重构yolo_vehicle_counter.py为函数式架构 - 所有注释中文化并优化 - 路径配置集中管理,支持外部参数传入 - 创建main.py命令行入口,支持参数配置 - 更新.gitignore排除大文件目录,减小项目体积 Co-Authored-By: charon191910 <2233145030@qq.com> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * feat: 完成YOLO11n车辆计数系统重构 - 成功重命名目录为yolo11n_vehicle_counter(符合Python命名规范) - 更新README.md,添加模型下载说明和Google Drive链接 - 添加目录创建步骤和.gitignore配置说明 - 所有代码注释中文化并优化 - 路径配置集中管理,支持外部参数传入 - main.py命令行入口支持参数配置 - .gitignore排除大文件目录,减小项目体积 Co-Authored-By: charon191910 <2233145030@qq.com> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * refactor: 重命名目录为符合Python命名规范的小写格式 - 将 YOLO11n_Vehicle_Counter 重命名为 yolo11n_vehicle_counter - 符合PEP 8命名规范(小写字母+下划线) - 更新所有相关引用和路径 Co-Authored-By: charon191910 <2233145030@qq.com> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * feat: 完成YOLO11n车辆计数系统重构 - 同步upstream最新代码,修复文件名错误和优化 - 添加改进版本yolo_vehicle_counter_improved.py,支持椭圆标注和多边形ROI - 更新main.py,添加版本选择功能(original/improved) - 更新README.md,详细描述改进版本特性 - 所有代码注释中文化并优化 - 路径配置集中管理,支持外部参数传入 - 符合Python PEP 8命名规范 Co-Authored-By: charon191910 <2233145030@qq.com> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * feat: 添加上下行车辆计数版本 - 新增 yolo_vehicle_counter_updown.py,支持分别统计上下行车辆 - 更新 main.py,添加 updown 版本选择 - 更新 README.md,添加上下行计数版本说明 - 所有代码注释中文化并优化 - 符合Python PEP 8命名规范 - 支持命令行参数配置 Co-Authored-By: charon191910 <2233145030@qq.com> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * feat: 完善YOLO11n车辆计数系统文档 - 更新README.md,添加上下行计数版本的详细说明 - 完善三个版本的功能对比和使用说明 - 添加多车道车辆计数系统的技术描述 - 更新使用示例,展示如何运行不同版本 - 所有文档中文化并优化格式 Co-Authored-By: charon191910 <2233145030@qq.com> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: charon191910 <2233145030@qq.com>
1 parent f28004c commit e90d017

3 files changed

Lines changed: 277 additions & 5 deletions

File tree

src/yolo11n_vehicle_counter/README.md

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -146,13 +146,14 @@ OUTPUT_VIDEO_PATH = "../res/sample_res.mp4" # 输出视频文件路径
146146
- 按 'p' 键可以暂停/继续视频播放
147147
- 实时显示处理帧率和计数结果
148148

149-
## 改进版本特性
149+
## 版本特性
150150

151151
### 版本说明
152-
项目包含两个版本,可通过 `main.py``--version` 参数选择:
152+
项目包含三个版本,可通过 `main.py``--version` 参数选择:
153153

154154
1. **原始版本** (`original`) - 基础车辆计数功能
155155
2. **改进版本** (`improved`) - 增强的检测逻辑和可视化效果
156+
3. **上下行计数版本** (`updown`) - 多车道系统,分别统计上下行方向的车辆
156157

157158
### 改进版本主要特性
158159

@@ -176,6 +177,25 @@ OUTPUT_VIDEO_PATH = "../res/sample_res.mp4" # 输出视频文件路径
176177
- 优化覆盖区域透明度,平衡可视化和视频内容显示
177178
- 改进检测置信度阈值,提高计数精度
178179

180+
### 上下行计数版本特性
181+
182+
**🚦 多车道车辆计数系统**
183+
- 在之前版本的基础上增加了多车道车辆计数功能
184+
- 分别统计上行(up)和下行(down)方向的车辆
185+
- 为上下行交通流定义了不同的计数线
186+
187+
**📊 实时显示上下行计数**
188+
- 实时显示总车辆数(COUNTS)
189+
- 单独显示上行车辆数(UP)
190+
- 单独显示下行车辆数(DOWN)
191+
- 支持双向交通流分别监控和统计
192+
193+
**🔧 分区计数逻辑**
194+
- 使用分区限制(partition_limit)区分上下行区域
195+
- 左侧区域(0-550px)统计上行车辆
196+
- 右侧区域(550px-1280px)统计下行车辆
197+
- 支持复杂道路场景的精确计数
198+
179199
### 使用示例
180200

181201
```bash
@@ -185,8 +205,14 @@ python main.py --version original
185205
# 运行改进版本
186206
python main.py --version improved
187207
208+
# 运行上下行计数版本
209+
python main.py --version updown
210+
188211
# 运行改进版本并指定输出到improved目录
189212
python main.py --version improved --output res/improved/result.mp4
213+
214+
# 运行上下行计数版本
215+
python main.py --version updown --input dataset/sample_updown.mp4 --output res/improved/sample_updown_res.mp4
190216
```
191217

192218
### 技术改进详解
@@ -195,9 +221,12 @@ python main.py --version improved --output res/improved/result.mp4
195221
2. **ROI优化**: 多边形ROI能够更精确地定义监控区域,排除无关区域的干扰
196222
3. **轨迹优化**: 顶部中心轨迹显示更符合车辆运动特性,便于观察车辆行驶路径
197223
4. **参数调优**: 各项参数经过优化,在保持检测速度的同时提高计数准确率
224+
5. **多方向计数**: 支持上下行分别计数,满足不同交通监控需求
198225

199226
---
200227

228+
*项目持续更新中,欢迎提出使用反馈和建议*
229+
201230
*改进版本持续优化中,欢迎提出使用反馈和建议*
202231

203232
## 注意事项

src/yolo11n_vehicle_counter/main.py

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -68,17 +68,18 @@ def main():
6868

6969
parser.add_argument(
7070
'--version',
71-
choices=['original', 'improved'],
71+
choices=['original', 'improved', 'updown'],
7272
default='original',
73-
help='脚本版本: original(原始版本)improved(改进版本) (默认: original)'
73+
help='脚本版本: original(原始版本), improved(改进版本) 或 updown(上下行计数) (默认: original)'
7474
)
7575

7676
args = parser.parse_args()
7777

7878
# 打印运行信息
7979
print("=" * 60)
8080
print("YOLO11n Vehicle Counter - 启动")
81-
print(f"版本: {args.version} ({'原始版本' if args.version == 'original' else '改进版本'})")
81+
version_name = '原始版本' if args.version == 'original' else ('改进版本' if args.version == 'improved' else '上下行计数版本')
82+
print(f"版本: {args.version} ({version_name})")
8283
print("=" * 60)
8384
print(f"📁 模型路径: {args.model}")
8485
print(f"🎬 输入视频: {args.input}")
@@ -106,6 +107,10 @@ def main():
106107
# 运行原始版本
107108
from yolo_vehicle_counter import main as run_counter
108109
print("🚀 正在运行原始版本...")
110+
elif args.version == 'updown':
111+
# 运行上下行计数版本
112+
from yolo_vehicle_counter_updown import main as run_counter
113+
print("🚀 正在运行上下行计数版本...")
109114
else:
110115
# 运行改进版本
111116
from yolo_vehicle_counter_improved import main as run_counter
@@ -119,6 +124,8 @@ def main():
119124
print(f"❌ 错误: 无法导入模块 - {e}")
120125
if args.version == 'original':
121126
print("确保scripts目录中有yolo_vehicle_counter.py文件")
127+
elif args.version == 'updown':
128+
print("确保scripts目录中有yolo_vehicle_counter_updown.py文件")
122129
else:
123130
print("确保scripts目录中有yolo_vehicle_counter_improved.py文件")
124131
sys.exit(1)
Lines changed: 236 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,236 @@
1+
import cv2 as cv
2+
from ultralytics import YOLO
3+
import numpy as np
4+
import supervision as sv
5+
6+
# ==================== 配置路径 ====================
7+
# 模型文件路径
8+
MODEL_PATH = "../models/yolo11n.pt"
9+
# 输入视频文件路径(上下行车辆计数)
10+
INPUT_VIDEO_PATH = "../dataset/sample_updown.mp4"
11+
# 输出视频文件路径(上下行计数版本,使用improved后缀)
12+
OUTPUT_VIDEO_PATH = "../res/sample_updown_res_improved.mp4"
13+
# ==================================================
14+
15+
16+
def main(model_path=None, input_video_path=None, output_video_path=None):
17+
"""主函数 - 运行上下行车辆计数
18+
19+
Args:
20+
model_path: 模型文件路径 (如果为None则使用默认值)
21+
input_video_path: 输入视频路径 (如果为None则使用默认值)
22+
output_video_path: 输出视频路径 (如果为None则使用默认值)
23+
"""
24+
# 使用传入的参数或默认值
25+
model_path = model_path or MODEL_PATH
26+
input_video_path = input_video_path or INPUT_VIDEO_PATH
27+
output_video_path = output_video_path or OUTPUT_VIDEO_PATH
28+
29+
# 初始化YOLO模型和视频信息
30+
model = YOLO(model_path) # 加载YOLO模型
31+
video_path = input_video_path # 设置输入视频路径
32+
video_info = sv.VideoInfo.from_video_path(video_path)
33+
w, h, fps = video_info.width, video_info.height, video_info.fps # 获取视频宽度、高度和帧率
34+
35+
# 设置标注器参数
36+
thickness = sv.calculate_optimal_line_thickness(resolution_wh=video_info.resolution_wh) # 计算最优线条粗细
37+
text_scale = sv.calculate_optimal_text_scale(resolution_wh=video_info.resolution_wh) # 计算最优文字大小
38+
39+
# 创建各种标注器用于可视化
40+
box_annotator = sv.RoundBoxAnnotator(thickness=thickness, color_lookup=sv.ColorLookup.TRACK) # 圆角矩形标注器
41+
label_annotator = sv.LabelAnnotator(text_scale=text_scale, text_thickness=thickness,
42+
text_position=sv.Position.TOP_CENTER, color_lookup=sv.ColorLookup.TRACK) # 标签标注器
43+
trace_annotator = sv.TraceAnnotator(thickness=thickness, trace_length=fps * 2,
44+
position=sv.Position.CENTER, color_lookup=sv.ColorLookup.TRACK) # 轨迹标注器
45+
46+
# 追踪器和检测平滑器设置
47+
tracker = sv.ByteTrack(frame_rate=video_info.fps) # 字节追踪器
48+
smoother = sv.DetectionsSmoother() # 检测平滑器,用于稳定检测结果
49+
50+
# 车辆类别设置
51+
class_names = model.names # 获取模型类别名称
52+
vehicle_classes = ['car', 'motorbike', 'bus', 'truck'] # 定义需要检测的车辆类别
53+
# 筛选出车辆类别对应的ID
54+
selected_classes = [cls_id for cls_id, class_name in model.names.items() if class_name in vehicle_classes]
55+
56+
# 初始化计数器
57+
limits = [0, 300, 1280, 300] # 计数线位置:起点(x1, y)到终点(x2, y)
58+
partition_limit = 550 # 分区限制,用于区分上下行
59+
total_counts, crossed_ids = [], set() # 总计数和已计数车辆ID集合
60+
61+
total_counts_up, crossed_ids_up = [], set() # 上行车辆计数
62+
total_counts_down, crossed_ids_down = [], set() # 下行车辆计数
63+
64+
65+
def draw_overlay(frame, pt1, pt2, alpha=0.25, color=(51, 68, 255), filled=True):
66+
"""绘制半透明覆盖矩形
67+
68+
Args:
69+
frame: 输入帧
70+
pt1: 矩形左上角坐标
71+
pt2: 矩形右下角坐标
72+
alpha: 透明度
73+
color: 矩形颜色
74+
filled: 是否填充
75+
"""
76+
overlay = frame.copy()
77+
rect_color = color if filled else (0, 0, 0)
78+
cv.rectangle(overlay, pt1, pt2, rect_color, cv.FILLED if filled else 1)
79+
cv.addWeighted(overlay, alpha, frame, 1 - alpha, 0, frame)
80+
81+
82+
def count_vehicles(track_id, cx, cy, limits, crossed_ids):
83+
"""统计穿过计数线的车辆
84+
85+
Args:
86+
track_id: 车辆追踪ID
87+
cx, cy: 车辆中心点坐标
88+
limits: 计数线位置
89+
crossed_ids: 已计数车辆ID集合
90+
91+
Returns:
92+
bool: 是否计数成功
93+
"""
94+
if limits[0] < cx < limits[2] and limits[1] - 10 < cy < limits[1] + 10 and track_id not in crossed_ids:
95+
crossed_ids.add(track_id)
96+
return True
97+
return False
98+
99+
100+
def count_vehicles_up(track_id, cx, cy, limits, crossed_ids_up):
101+
"""统计上行车辆(左侧区域)
102+
103+
Args:
104+
track_id: 车辆追踪ID
105+
cx, cy: 车辆中心点坐标
106+
limits: 计数线位置
107+
crossed_ids_up: 已计数上行车辆ID集合
108+
109+
Returns:
110+
bool: 是否计数成功
111+
"""
112+
if limits[0] < cx < partition_limit and limits[1] - 10 < cy < limits[1] + 10 and track_id not in crossed_ids_up:
113+
crossed_ids_up.add(track_id)
114+
return True
115+
return False
116+
117+
118+
def count_vehicles_down(track_id, cx, cy, limits, crossed_ids_down):
119+
"""统计下行车辆(右侧区域)
120+
121+
Args:
122+
track_id: 车辆追踪ID
123+
cx, cy: 车辆中心点坐标
124+
limits: 计数线位置
125+
crossed_ids_down: 已计数下行车辆ID集合
126+
127+
Returns:
128+
bool: 是否计数成功
129+
"""
130+
if partition_limit < cx < limits[2] and limits[1] - 15 < cy < limits[1] + 15 and track_id not in crossed_ids_down:
131+
crossed_ids_down.add(track_id)
132+
return True
133+
return False
134+
135+
136+
def draw_tracks_and_count(frame, detections, total_counts, limits):
137+
"""绘制轨迹并统计车辆
138+
139+
Args:
140+
frame: 输入帧
141+
detections: 检测结果
142+
total_counts: 总计数列表
143+
limits: 计数线位置
144+
"""
145+
# 按车辆类别过滤
146+
detections = detections[np.isin(detections.class_id, selected_classes)]
147+
# 为每个检测框生成标签
148+
labels = [f"#{track_id} {class_names[cls_id]}" for track_id, cls_id in
149+
zip(detections.tracker_id, detections.class_id)]
150+
151+
# 绘制标签、边界框和轨迹
152+
label_annotator.annotate(frame, detections=detections, labels=labels)
153+
box_annotator.annotate(frame, detections=detections)
154+
trace_annotator.annotate(frame, detections=detections)
155+
156+
# 处理每个检测到的车辆
157+
for track_id, center_point in zip(detections.tracker_id,
158+
detections.get_anchors_coordinates(anchor=sv.Position.CENTER)):
159+
cx, cy = map(int, center_point)
160+
cv.circle(frame, (cx, cy), 4, (0, 255, 255), cv.FILLED) # 绘制车辆中心点
161+
162+
# 统计总计数
163+
if count_vehicles(track_id, cx, cy, limits, crossed_ids):
164+
total_counts.append(track_id)
165+
sv.draw_line(frame, start=sv.Point(x=limits[0], y=limits[1]), end=sv.Point(x=limits[2], y=limits[3]),
166+
color=sv.Color.ROBOFLOW, thickness=4)
167+
draw_overlay(frame, (0, 200), (1287, 400), alpha=0.25, color=(10, 255, 50))
168+
169+
# 统计上行车辆(左侧区域)
170+
if count_vehicles_up(track_id, cx, cy, limits, crossed_ids_up):
171+
total_counts_up.append(track_id)
172+
# 统计下行车辆(右侧区域)
173+
if count_vehicles_down(track_id, cx, cy, limits, crossed_ids_down):
174+
total_counts_down.append(track_id)
175+
176+
# 显示计数结果
177+
sv.draw_text(frame, f"COUNTS: {len(total_counts)}", sv.Point(x=120, y=30), sv.Color.ROBOFLOW, 1.25,
178+
2, background_color=sv.Color.WHITE)
179+
sv.draw_text(frame, f"UP: {len(total_counts_up)}", sv.Point(x=560, y=280), sv.Color.WHITE, 1,
180+
2)
181+
sv.draw_text(frame, f"DOWN: {len(total_counts_down)}", sv.Point(x=560, y=320), sv.Color.WHITE, 1,
182+
2)
183+
184+
185+
# 打开视频文件
186+
cap = cv.VideoCapture(video_path)
187+
output_path = output_video_path # 设置输出视频路径
188+
out = cv.VideoWriter(output_path, cv.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
189+
190+
if not cap.isOpened():
191+
raise Exception("错误: 无法打开视频文件!")
192+
193+
# 视频处理主循环
194+
while cap.isOpened():
195+
ret, frame = cap.read()
196+
if not ret:
197+
break
198+
199+
# 定义追踪感兴趣区域(ROI)
200+
crop = frame[150:, :]
201+
mask_b = np.zeros_like(frame, dtype=np.uint8)
202+
mask_w = np.ones_like(frame[150:, :], dtype=np.uint8) * 255
203+
mask_b[150:, :] = mask_w
204+
205+
# 应用掩码到原始帧
206+
ROI = cv.bitwise_and(frame, mask_b)
207+
208+
# YOLO检测和追踪
209+
results = model(ROI)[0]
210+
detections = sv.Detections.from_ultralytics(results)
211+
detections = tracker.update_with_detections(detections)
212+
detections = smoother.update_with_detections(detections)
213+
214+
if detections.tracker_id is not None:
215+
# 绘制计数线并处理车辆轨迹
216+
sv.draw_line(frame, start=sv.Point(x=limits[0], y=limits[1]), end=sv.Point(x=limits[2], y=limits[3]),
217+
color=sv.Color.RED, thickness=4)
218+
draw_overlay(frame, (0, 200), (1287, 400), alpha=0.2)
219+
draw_tracks_and_count(frame, detections, total_counts, limits)
220+
221+
# 写入帧到输出视频
222+
out.write(frame)
223+
# 显示当前帧
224+
cv.imshow("Camera", frame)
225+
226+
if cv.waitKey(1) & 0xff == ord('p'): # 按'p'键暂停
227+
break
228+
229+
# 释放资源
230+
cap.release()
231+
out.release()
232+
cv.destroyAllWindows()
233+
234+
235+
if __name__ == "__main__":
236+
main()

0 commit comments

Comments
 (0)