init repo

.gitignore

@@ -0,0 +1,10 @@
+__pycache__
+
+runs
+
+*.mp4
+*.dll
+
+*.png
+
+*.pt

README.md

@@ -0,0 +1,20 @@
+# YOLOv8 过线检测
+
+基于 YOLOv8 的过线检测，支持自定义划线。
+
+## 开始
+
+建议的 Python 版本为 3.10.x，首先创建一个虚拟环境（推荐），然后按顺序安装以下包
+
+```shell
+# GPU 环境
+pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
+pip install ultralytics
+pip install supervision opencv-contrib-python pillow lapx
+```
+
+然后可以开始运行
+
+```shell
+python predictor.py
+```shell

line_selector.py

@@ -0,0 +1,66 @@
+import tkinter as tk
+
+class LineSelector:
+    def __init__(self, image_path):
+        self.image_path = image_path
+        self.start_x, self.start_y = None, None
+        self.end_x, self.end_y = None, None
+        self.line_id = None
+        self.coordinates = None
+        self.success = False
+
+    def on_mouse_press(self, event):
+        self.start_x, self.start_y = event.x, event.y
+
+    def on_mouse_drag(self, event):
+        if self.line_id:
+            self.canvas.delete(self.line_id)
+        self.end_x, self.end_y = event.x, event.y
+        self.line_id = self.canvas.create_line(self.start_x, self.start_y, self.end_x, self.end_y, fill="white")
+
+    def on_mouse_release(self, event):
+        self.end_x, self.end_y = event.x, event.y
+        self.update_coordinates_label()
+
+    def update_coordinates_label(self):
+        coordinates = f"起点坐标: ({self.start_x}, {self.start_y})，终点坐标: ({self.end_x}, {self.end_y})"
+        self.coordinates_label.config(text=coordinates)
+    
+    def on_confirm_button_click(self):
+        self.success = True
+        self.root.destroy()
+        self.root.quit()
+
+    def draw_image_and_get_coordinates(self):
+        self.root = tk.Tk()
+        self.root.title("绘制检测线")
+        self.canvas = tk.Canvas(self.root, width=1920, height=1080)
+        self.canvas.pack()
+        image = tk.PhotoImage(file=self.image_path)
+        self.canvas.create_image(0, 0, anchor=tk.NW, image=image)
+
+        self.coordinates_label = tk.Label(self.root, text="", fg="red")
+        self.coordinates_label.pack()
+
+        confirm_button = tk.Button(self.root, text="确认", command=self.on_confirm_button_click)
+        confirm_button.pack()
+
+        self.canvas.bind("<ButtonPress-1>", self.on_mouse_press)
+        self.canvas.bind("<B1-Motion>", self.on_mouse_drag)
+        self.canvas.bind("<ButtonRelease-1>", self.on_mouse_release)
+
+        self.root.mainloop()
+
+    def get_coordinates(self):
+        self.coordinates = [self.start_x, self.start_y, self.end_x, self.end_y]
+        return self.coordinates
+
+def main():
+    image_path = "background.png"
+    selector = LineSelector(image_path)
+    selector.draw_image_and_get_coordinates()
+    coordinates = selector.get_coordinates()
+    print("坐标:", coordinates)
+
+if __name__ == "__main__":
+    main()

predictor.py

@@ -0,0 +1,170 @@
+# YOLOv8
+from ultralytics import YOLO
+
+# supervision
+import supervision as sv
+from supervision.draw.color import Color
+
+# OpenCV
+import cv2
+
+# Line Selector
+from line_selector import LineSelector
+
+# time date
+from datetime import datetime
+import time
+
+# 视频源
+source = "line.mp4"
+source = "rtmp://10.0.0.21:1935/live/picam3"
+
+# 读取背景图像
+cap = cv2.VideoCapture(source)
+
+fps = cap.get(cv2.CAP_PROP_FPS)
+width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+
+# 输出视频
+
+########## 配置是否录制 ##########
+record = True
+
+if record:
+    output = "output2.mp4"
+    out = cv2.VideoWriter(output, cv2.VideoWriter_fourcc(*"avc1"), fps / 2, (width, height))
+    out.set(cv2.VIDEOWRITER_PROP_QUALITY, 80)
+
+# 读取背景图像
+ret, frame = cap.read()
+if ret:
+    cv2.imwrite("background.png", frame)
+    selector = LineSelector("background.png")
+    selector.draw_image_and_get_coordinates()
+    if selector.success:
+        coordinates = selector.get_coordinates()
+        print("线段坐标: ", coordinates)
+    else:
+        print("未选择线段坐标")
+        cap.release()
+        exit(0)
+
+cap.release()
+
+# 加载目标检测模型
+print("正在加载模型...")
+model = YOLO("yolov8x.pt")
+
+# 越线检测位置
+LINE_START = sv.Point(coordinates[0], coordinates[1])
+LINE_END = sv.Point(coordinates[2], coordinates[3])
+line_counter = sv.LineZone(start=LINE_START, end=LINE_END)
+
+# 线的可视化配置
+line_color = Color(r=224, g=57, b=151)
+line_annotator = sv.LineZoneAnnotator(
+    thickness=2, text_thickness=2, text_scale=1, color=line_color, text_offset=2.0, custom_in_text="North", custom_out_text="South"
+)
+
+# 目标检测可视化配置
+box_annotator = sv.BoxAnnotator(thickness=2, text_thickness=2, text_scale=1)
+
+# 起始时间
+start_time = time.perf_counter()
+
+# 限定帧数
+########## 配置是否限定帧数 ##########
+frame_limit = False
+frame_limit_upper = 43000
+frame_count = 0
+
+# 逐帧跟踪
+for result in model.track(source, device=0, verbose=False, stream=True):
+    # 限定帧数
+    frame_count += 1
+    if frame_count == frame_limit_upper and frame_limit:
+        print(f"已经到达{frame_limit_upper}帧，停止运行")
+        break
+
+    # 获取原始图像
+    frame = result.orig_img
+
+    # 用 supervision 解析预测结果
+    detections = sv.Detections.from_ultralytics(result)
+
+    ## 过滤掉某些类别
+    # detections = detections[(detections.class_id != 60) & (detections.class_id != 0)]
+
+    # 解析追踪ID
+    if result.boxes.id is None:
+        continue
+
+    detections.tracker_id = result.boxes.id.numpy().astype(int)
+
+    # 获取每个目标的：追踪ID、类别名称、置信度
+    class_ids = detections.class_id  # 类别ID
+    confidences = detections.confidence  # 置信度
+    tracker_ids = detections.tracker_id  # 多目标追踪ID
+    labels = [
+        "#{} {} {:.1f}".format(
+            tracker_ids[i], model.names[class_ids[i]], confidences[i] * 100
+        )
+        for i in range(len(class_ids))
+    ]
+
+    # 绘制目标检测可视化结果
+    frame = box_annotator.annotate(scene=frame, detections=detections, labels=labels)
+
+    # 越线检测
+    line_counter.trigger(detections=detections)
+    line_annotator.annotate(frame=frame, line_counter=line_counter)
+
+    # 显示日期
+    cv2.putText(
+        frame,
+        datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+        (10, 30),
+        cv2.FONT_HERSHEY_SIMPLEX,
+        1,
+        (255, 255, 255),
+        2,
+        cv2.LINE_AA,
+    )
+
+    # 结束时间
+    end_time = time.perf_counter()
+
+    # 计算帧率
+    fps = 1.0 / (end_time - start_time)
+
+    # 显示帧率
+
+    cv2.putText(
+        frame,
+        "FPS: {:.2f}".format(fps),
+        (10, 60),
+        cv2.FONT_HERSHEY_SIMPLEX,
+        1,
+        (255, 255, 255),
+        2,
+        cv2.LINE_AA,
+    )
+
+    # 保存视频
+    if record:
+        out.write(frame)
+
+    # 显示结果
+    cv2.imshow("Frame", frame)
+
+    if cv2.waitKey(1) & 0xFF == ord("q"):
+        break
+
+    # 起始时间
+    start_time = time.perf_counter()
+
+# 释放资源
+if record:
+    out.release()
+cv2.destroyAllWindows()