Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 62 additions & 28 deletions src/yolo11n_vehicle_counter/scripts/yolo_vehicle_counter_carla.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,15 +57,17 @@ def main(model_path=None, input_video_path=None, output_video_path=None):
selected_classes = [cls_id for cls_id, class_name in model.names.items() if class_name in vehicle_classes]

# 初始化计数器
# 针对CARLA视频优化:红线再往右移,更容易捕捉到车辆
# 原limits: [400, 400, 1250, 400] - 红线在y=400,太高了
# 新limits: [350, 500, 1230, 500] - 红线在y=500,进一步往右移动
limits = [350, 500, 1230, 500] # 计数线位置:起点(x1, y)到终点(x2, y) - 针对CARLA视频放低并进一步往右移红线
limits = [350, 750, 1230, 750] # 计数线位置
total_counts, crossed_ids = [], set() # 总计数和已计数车辆ID集合
track_history = {} # 轨迹历史: {track_id: [y1, y2, ...]}
counted_tracks = set() # 已完成计数的轨迹ID

# 分类计数器
class_counts = {'car': 0, 'motorbike': 0, 'bus': 0, 'truck': 0} # 各类别已计数ID集合
crossed_by_class = {cls: set() for cls in class_counts.keys()} # 各类别已计数ID集合

# 轨迹历史记录(用于判断穿越方向)
previous_y = {} # 上一帧的y坐标


def draw_overlay(frame, pt1, pt2, alpha=0.25, color=(51, 68, 255), filled=True):
Expand Down Expand Up @@ -120,8 +122,8 @@ def draw_tracks_and_count(frame, detections, limits):
"""
nonlocal total_counts, crossed_ids, class_counts, crossed_by_class

# 按车辆类别和检测置信度过滤 - 针对CARLA视频优化置信度阈值
detections = detections[(np.isin(detections.class_id, selected_classes)) & (detections.confidence > 0.4)]
# 按车辆类别和检测置信度过滤 - 降低阈值提高检测率
detections = detections[(np.isin(detections.class_id, selected_classes)) & (detections.confidence > 0.15)]

# 为每个检测框生成标签
labels = [f"#{track_id} {class_names[cls_id]}" for track_id, cls_id in
Expand All @@ -141,19 +143,40 @@ def draw_tracks_and_count(frame, detections, limits):

cv.circle(frame, (cx, cy), 4, (0, 255, 255), cv.FILLED) # 绘制车辆中心点

# 检查是否穿过计数线
if limits[0] < cx < limits[2] and limits[1] - 15 < cy < limits[1] + 15:
if track_id not in crossed_ids:
crossed_ids.add(track_id)
total_counts.append(track_id)

# 分类计数
if cls_name in class_counts:
class_counts[cls_name] += 1

sv.draw_line(frame, start=sv.Point(x=limits[0], y=limits[1]), end=sv.Point(x=limits[2], y=limits[3]),
color=sv.Color.ROBOFLOW, thickness=4)
draw_overlay(frame, (350, 450), (1230, 550), alpha=0.25, color=(10, 255, 50))
# 记录轨迹历史
if track_id not in track_history:
track_history[track_id] = []
previous_y[track_id] = cy # 初始化上一帧y坐标
track_history[track_id].append(cy)
# 只保留最近30帧历史
if len(track_history[track_id]) > 30:
track_history[track_id] = track_history[track_id][-30:]

# 严谨的穿越方向判断
if track_id not in crossed_ids:
prev_y = previous_y.get(track_id, cy)
curr_y = cy
line_y = limits[1]

# 检查是否在线的x范围内
if limits[0] < cx < limits[2]:
# 判断是否从上往下穿越计数线
crossed_line = (prev_y < line_y <= curr_y) or (prev_y < line_y and curr_y > prev_y + 10)
if crossed_line:
crossed_ids.add(track_id)
counted_tracks.add(track_id)
total_counts.append(track_id)

# 分类计数
if cls_name in class_counts:
class_counts[cls_name] += 1

sv.draw_line(frame, start=sv.Point(x=limits[0], y=limits[1]), end=sv.Point(x=limits[2], y=limits[3]),
color=sv.Color.ROBOFLOW, thickness=4)
draw_overlay(frame, (350, 700), (1230, 800), alpha=0.25, color=(10, 255, 50))

# 更新上一帧y坐标
previous_y[track_id] = curr_y

# 显示车辆总计数
sv.draw_text(frame, f"TOTAL: {len(total_counts)}", sv.Point(x=50, y=50), sv.Color.ROBOFLOW, 0.8,
Expand All @@ -179,11 +202,10 @@ def draw_tracks_and_count(frame, detections, limits):
if frame_count % 30 == 0: # 每30帧打印一次进度
print(f"处理进度: 第 {frame_count} 帧, 已计数: {len(total_counts)} 辆车")

# 针对CARLA视频优化ROI区域 - 减少不必要的区域
# CARLA视频通常上半部分有天空,下半部分有车辆
# 调整ROI为中间到下半部分,提高检测效率
roi_top = h // 3 # 从1/3高度开始
roi_left = w // 6 # 从1/6宽度开始
# 针对CARLA视频优化ROI区域 - 扩大检测范围避免漏检
# 只裁剪左右边缘和最顶部天空区域,保留更多车辆检测机会
roi_top = h // 6 # 从1/6高度开始(之前是1/3,丢失太多)
roi_left = w // 12 # 从1/12宽度开始(之前是1/6,减少裁剪)
crop = frame[roi_top:, roi_left:w-roi_left]
mask_b = np.zeros_like(frame, dtype=np.uint8)
mask_w = np.ones_like(crop, dtype=np.uint8) * 255
Expand All @@ -192,9 +214,21 @@ def draw_tracks_and_count(frame, detections, limits):
# 应用掩码到原始帧
ROI = cv.bitwise_and(frame, mask_b)

# YOLO检测和追踪
results = model(ROI)[0]
detections = sv.Detections.from_ultralytics(results)
# YOLO多尺度检测 - 使用不同输入尺寸提高远处车辆检测率
# 原始尺寸用于近处大目标,640尺寸用于远处小目标
results_main = model(ROI, imgsz=1280)[0] # 增大输入尺寸,检测更多细节
detections_main = sv.Detections.from_ultralytics(results_main)

# 如果需要额外补充检测(可选)
if len(detections_main) < 3: # 检测过少时尝试小目标检测
results_small = model(ROI, imgsz=1920)[0]
detections_small = sv.Detections.from_ultralytics(results_small)
# 合并结果(去重)
all_detections = detections_main if len(detections_main) > 0 else detections_small
else:
all_detections = detections_main

detections = all_detections
detections = tracker.update_with_detections(detections)
detections = smoother.update_with_detections(detections)

Expand All @@ -203,7 +237,7 @@ def draw_tracks_and_count(frame, detections, limits):
sv.draw_line(frame, start=sv.Point(x=limits[0], y=limits[1]), end=sv.Point(x=limits[2], y=limits[3]),
color=sv.Color.RED, thickness=4)
# 调整覆盖区域透明度 - 与红线位置匹配
draw_overlay(frame, (350, 450), (1230, 550), alpha=0.15)
draw_overlay(frame, (350, 700), (1230, 800), alpha=0.15)
draw_tracks_and_count(frame, detections, limits)

# 写入帧到输出视频
Expand Down
Loading