Ultralytics YOLOv8.0.226 进出计数

in_and_out_counting

对本地视频文件进行出入计数

from ultralytics import YOLO
from collections import defaultdict
import cv2
from ultralytics.utils.checks import check_imshow, check_requirements
from ultralytics.utils.plotting import Annotator, colors
from shapely.geometry import Polygon
from shapely.geometry.point import Point
from datetime import datetime, timedelta
import re
from ultralytics.utils import LOGGER, colorstr
from openpyxl import Workbook

check_requirements('shapely>=2.0.0')


class ObjectCounter:

    def __init__(self):
        self.is_drawing = False
        self.selected_point = None
        self.reg_pts = None
        self.counting_region = None
        self.region_color = (255, 255, 255)
        self.im0 = None
        self.tf = None
        self.view_img = True
        self.names = None
        self.annotator = None
        self.in_counts = 0
        self.out_counts = 0
        self.counting_list = []
        self.track_history = defaultdict(list)
        self.track_thickness = 2
        self.draw_tracks = True
        self.env_check = check_imshow(warn=True)
        self.incount_label = 'InCount: 0'
        self.outcount_label = 'OutCount: 0'
        self.k = 1

    def set_args(self, classes_names, reg_pts, region_color=None, line_thickness=2, track_thickness=2, view_img=True, draw_tracks=True):
        self.tf = line_thickness
        self.view_img = view_img
        self.track_thickness = track_thickness
        self.draw_tracks = draw_tracks
        self.reg_pts = reg_pts
        self.counting_region = Polygon(self.reg_pts)
        self.names = classes_names
        self.region_color = region_color if region_color else self.region_color

    def mouse_event_for_region(self, event, x, y, flags, params):
        if event == cv2.EVENT_LBUTTONDOWN:
            for i, point in enumerate(self.reg_pts):
                if isinstance(point, (tuple, list)) and len(point) >= 2:
                    if abs(x - point[0]) < 10 and abs(y - point[1]) < 10:
                        self.selected_point = i
                        self.is_drawing = True
                        break
        elif event == cv2.EVENT_MOUSEMOVE:
            if self.is_drawing and self.selected_point is not None:
                self.reg_pts[self.selected_point] = (x, y)
                self.counting_region = Polygon(self.reg_pts)
        elif event == cv2.EVENT_LBUTTONUP:
            self.is_drawing = False
            self.selected_point = None

    def extract_and_process_tracks(self, tracks):
        boxes = tracks[0].boxes.xyxy.cpu()
        clss = tracks[0].boxes.cls.cpu().tolist()
        track_ids = tracks[0].boxes.id.int().cpu().tolist()