二、项目完整代码
帮我们将这个代码，复制粘贴到我们YOLOv8的仓库里然后创建一个py文件存放进去即可。
python
import argparse
from collections import defaultdict
from pathlib import Path
import cv2
import numpy as np
from shapely.geometry import Polygon
from shapely.geometry.point import Point
from ultralytics import YOLO
from ultralytics.utils.files import increment_path
from ultralytics.utils.plotting import Annotator, colors
track_history = defaultdict(list)
current_region = None
counting_regions = [
{
'name': 'YOLOv8 Polygon Region',
'polygon': Polygon([(50, 80), (250, 20), (450, 80), (400, 350), (100, 350)]), # Polygon points
'counts': 0,
'dragging': False,
'region_color': (255, 42, 4), # BGR Value
'text_color': (255, 255, 255) # Region Text Color
},
{
'name': 'YOLOv8 Rectangle Region',
'polygon': Polygon([(200, 250), (440, 250), (440, 550), (200, 550)]), # Polygon points
'counts': 0,
'dragging': False,
'region_color': (37, 255, 225), # BGR Value
'text_color': (0, 0, 0), # Region Text Color
}, ]
def mouse_callback(event, x, y, flags, param):
"""Mouse call back event."""
global current_region
# Mouse left button down event
if event == cv2.EVENT_LBUTTONDOWN:
for region in counting_regions:
if region['polygon'].contains(Point((x, y))):
current_region = region
current_region['dragging'] = True
current_region['offset_x'] = x
current_region['offset_y'] = y
# Mouse move event
elif event == cv2.EVENT_MOUSEMOVE:
if current_region is not None and current_region['dragging']:
dx = x - current_region['offset_x']
dy = y - current_region['offset_y']
current_region['polygon'] = Polygon([
(p[0] + dx, p[1] + dy) for p in current_region['polygon'].exterior.coords])
current_region['offset_x'] = x
current_region['offset_y'] = y
# Mouse left button up event
elif event == cv2.EVENT_LBUTTONUP:
if current_region is not None and current_region['dragging']:
current_region['dragging'] = False
def run(
weights='yolov8n.pt',
source=None,
device='cpu',
view_img=False,
save_img=False,
exist_ok=False,
classes=None,
line_thickness=2,
track_thickness=2,
region_thickness=2,
):
"""
Run Region counting on a video using YOLOv8 and ByteTrack.
Supports movable region for real time counting inside specific area.
Supports multiple regions counting.
Regions can be Polygons or rectangle in shape
Args:
weights (str): Model weights path.
source (str): Video file path.
device (str): processing device cpu, 0, 1
view_img (bool): Show results.
save_img (bool): Save results.
exist_ok (bool): Overwrite existing files.
classes (list): classes to detect and track
line_thickness (int): Bounding box thickness.
track_thickness (int): Tracking line thickness
region_thickness (int): Region thickness.
"""
vid_frame_count = 0
# Check source path
if not Path(source).exists():
raise FileNotFoundError(f"Source path '{source}' does not exist.")
# Setup Model
model = YOLO(f'{weights}')
model.to('cuda') if device == '0' else model.to('cpu')
# Extract classes names
names = model.model.names
# Video setup
videocapture = cv2.VideoCapture(source)
frame_width, frame_height = int(videocapture.get(3)), int(videocapture.get(4))
fps, fourcc = int(videocapture.get(5)), cv2.VideoWriter_fourcc(*'mp4v')
# Output setup
save_dir = increment_path(Path('ultralytics_rc_output') / 'exp', exist_ok)
save_dir.mkdir(parents=True, exist_ok=True)
video_writer = cv2.VideoWriter(str(save_dir / f'{Path(source).stem}.mp4'), fourcc, fps, (frame_width, frame_height))
# Iterate over video frames
while videocapture.isOpened():
success, frame = videocapture.read()
if not success:
break
vid_frame_count += 1
# Extract the results
results = model.track(frame, persist=True, classes=classes)
if results[0].boxes.id is not None:
boxes = results[0].boxes.xyxy.cpu()
track_ids = results[0].boxes.id.int().cpu().tolist()
clss = results[0].boxes.cls.cpu().tolist()
annotator = Annotator(frame, line_width=line_thickness, example=str(names))
for box, track_id, cls in zip(boxes, track_ids, clss):
annotator.box_label(box, str(names[cls]), color=colors(cls, True))
bbox_center = (box[0] + box[2]) / 2, (box[1] + box[3]) / 2 # Bbox center
track = track_history[track_id] # Tracking Lines plot
track.append((float(bbox_center[0]), float(bbox_center[1])))
if len(track) > 30:
track.pop(0)
points = np.hstack(track).astype(np.int32).reshape((-1, 1, 2))
cv2.polylines(frame, [points], isClosed=False, color=colors(cls, True), thickness=track_thickness)
# Check if detection inside region
for region in counting_regions:
if region['polygon'].contains(Point((bbox_center[0], bbox_center[1]))):
region['counts'] += 1
# Draw regions (Polygons/Rectangles)
for region in counting_regions:
region_label = str(region['counts'])
region_color = region['region_color']
region_text_color = region['text_color']
polygon_coords = np.array(region['polygon'].exterior.coords, dtype=np.int32)
centroid_x, centroid_y = int(region['polygon'].centroid.x), int(region['polygon'].centroid.y)
text_size, _ = cv2.getTextSize(region_label,
cv2.FONT_HERSHEY_SIMPLEX,
fontScale=0.7,
thickness=line_thickness)
text_x = centroid_x - text_size[0] // 2
text_y = centroid_y + text_size[1] // 2
cv2.rectangle(frame, (text_x - 5, text_y - text_size[1] - 5), (text_x + text_size[0] + 5, text_y + 5),
region_color, -1)
cv2.putText(frame, region_label, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.7, region_text_color,
line_thickness)
cv2.polylines(frame, [polygon_coords], isClosed=True, color=region_color, thickness=region_thickness)
if view_img:
if vid_frame_count == 1:
cv2.namedWindow('Ultralytics YOLOv8 Region Counter Movable')
cv2.setMouseCallback('Ultralytics YOLOv8 Region Counter Movable', mouse_callback)
cv2.imshow('Ultralytics YOLOv8 Region Counter Movable', frame)
if save_img:
video_writer.write(frame)
for region in counting_regions: # Reinitialize count for each region
region['counts'] = 0
if cv2.waitKey(1) & 0xFF == ord('q'):
break
del vid_frame_count
video_writer.release()
videocapture.release()
cv2.destroyAllWindows()
def parse_opt():
"""Parse command line arguments."""
parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default='yolov8n.pt', help='initial weights path')
parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--source', type=str, default='替换你视频地址', help='video file path')
parser.add_argument('--view-img', action='store_true',default=True , help='show results')
parser.add_argument('--save-img', action='store_true', default=True, help='save results')
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --classes 0, or --classes 0 2 3')
parser.add_argument('--line-thickness', type=int, default=2, help='bounding box thickness')
parser.add_argument('--track-thickness', type=int, default=2, help='Tracking line thickness')
parser.add_argument('--region-thickness', type=int, default=4, help='Region thickness')
return parser.parse_args()
def main(opt):
"""Main function."""
run(**vars(opt))
if __name__ == '__main__':
opt = parse_opt()
main(opt)三、参数解析
下面上面项目核心代码的参数解析，共有10个，能够起到作用的参数并不多。
| 参数名 | 参数类型 | 参数讲解 | |
|---|---|---|---|
| 0 | weights | str | 用于检测视频的权重文件地址（可以是你训练好的，也可以是官方提供的） |
| 1 | device | str | 设备的选择可以用GPU也可以用CPU |
| 2 | source | str | 视频文件的地址，因为是用于视频检测，大家有兴趣其实可以将其改为摄像头的实时检测。 |
| 3 | view-img | bool | 是否显示视频结果 ，就是它在控制台会输出结果，如果设置为True就显示图像结果 |
| 4 | save-img | bool | 是否保存检测的结果，文件会存放在同级目录下的新文件夹内 |
| 5 | exist-ok | bool | 保存文件的名字检测的，大家不用理会这个参数 |
| 6 | classes | int | 这个参数比较重要，比如你的权重文件训练文件有多个类别，如果你只想要看某个特定类别的检测结果，只需要输入它的数字编号即可。 |
| 7 | line-thickness | int | 边界框的厚度，用于在检测到的对象周围绘制边界线 |
| 8 | track-thickness | int | 跟踪线的厚度。 |
| 9 | region-thickness | int | 区域框的厚度。 |
四、项目的使用教程
4.1 步骤一
我们在Yolo仓库的目录下创建一个py文件将代码存放进去，如下图所示。
4.2 步骤二
我们按照参数解析部分的介绍填好大家的参数，主要配置的有两个一个就是权重文件地址另一个就是视频的地址。
4.3 步骤三
我们挺好之后运行文件即可，此时会弹出视频框，其中有两个region框是可以滑动的以此来确定想要检测的目标区域，其原理就是Opencv中的mask掩码的原理使用（Gif有点糊大家对付看一下，毕竟只是起到一个示范作用）。
