Ultralytics YOLOv8.0.226 区域计数

regional_counting

对本地视频文件进行区域计数

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
from datetime import datetime, timedelta
import re
from ultralytics.utils import LOGGER, colorstr
import openpyxl

track_history = defaultdict(list)
current_region = None


def mouse_callback(event, x, y, flags, param):
    global current_region
    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
    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
    elif event == cv2.EVENT_LBUTTONUP:
        if current_region is not None and current_region['dragging']:
            current_region['dragging'] = False


def main():
    global counting_regions
    counting_regions = [
        {
            'name': 'YOLOv8 Polygon Region',
            'polygon': Polygon([(x1, y1), (x1 + width1, y1), (x1 + width1, y1 + height1), (x1, y1 + height1)]),
            'counts': 0,
            'dragging': True,
            'region_color': (255, 0, 0),
            'text_color': (0, 0, 0)
        },
        {
            'name': 'YOLOv8 Rectangle Region',
            'polygon': Polygon([(x2, y2), (x2 + width2, y2), (x2 + width2, y2 + height2), (x2, y2 + height2)]),  # Polygon points
            'counts': 0,
            'dragging': True,
            'region_color': (0, 255, 225),
            'text_color': (0, 0, 0),
        },
        {
            'name': 'YOLOv8 Rectangle Region',
            'polygon': Polygon([(x3, y3), (x3 + width3, y3), (x3 + width3, y3 + height3), (x3, y3 + height3)]),  # Polygon points
            'counts': 0,
            'dragging': True,
            'region_color': (0, 255, 0),
            'text_color': (0, 0, 0),
        }
    ]
    vid_frame_count = 0
    i = 1
    if not Path(video_path).exists():
        raise FileNotFoundError(f"Source path '{source}' does not exist.")
    model = YOLO(f'{weights}')
    model.to('cuda') if device == '0' else model.to('cpu')
    names = model.model.names
    videocapture = cv2.VideoCapture(video_path)
    frame_width, frame_height = int(videocapture.get(3)), int(videocapture.get(4))
    fps, fourcc = int(videocapture.get(5)), cv2.VideoWriter_fourcc(*'mp4v')