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')
    save_dir = increment_path(Path('runs'), exist_ok)
    save_dir.mkdir(parents=True, exist_ok=True)
    video_writer = cv2.VideoWriter(str(save_dir / 'regional_counting.mp4'), fourcc, fps, (frame_width, frame_height))
    year, month, day, hours, minutes, seconds = map(int, re.split('[-:\s]+', video_start_time))
    start_time = datetime(year, month, day, hours, minutes, seconds)
    workbook = openpyxl.Workbook()
    worksheet = workbook.active
    while videocapture.isOpened():
        success, frame = videocapture.read()
        if not success:
            break
        vid_frame_count += 1
        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
                track = track_history[track_id]
                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)
                for region in counting_regions:
                    if region['polygon'].contains(Point((bbox_center[0], bbox_center[1]))):
                        region['counts'] += 1
        number_of_people = []
        now_time = start_time + timedelta(seconds=int(vid_frame_count / fps))
        cv2.putText(frame, 'Frame time: ' + str(now_time), (120, 40), cv2.FONT_HERSHEY_COMPLEX, 0.7, (0, 0, 255), line_thickness)
        for region in counting_regions:
            number_of_people.append(region['counts'])
            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 now_time == start_time + timedelta(seconds=i * record_interval_time):
            worksheet.append([str(now_time)] + number_of_people)
            if i == 1:
                number_of_regions = len(counting_regions)
                record_table = [[0 for x in range(4)] for y in range(number_of_regions)]
                for x in range(number_of_regions):
                    for y in range(4):
                        if y == 0 or y == 2:
                            record_table[x][y] = now_time
                        else:
                            record_table[x][y] = number_of_people[x]
            for x in range(number_of_regions):
                if number_of_people[x] > record_table[x][1]:
                    record_table[x][1] = number_of_people[x]
                    record_table[x][0] = now_time
                if number_of_people[x] < record_table[x][3]:
                    record_table[x][3] = number_of_people[x]
                    record_table[x][2] = now_time
            i = i + 1
            detection_information = str(colorstr('magenta', "Time: ")) + str(colorstr('red', now_time))
            for x in range(number_of_regions):
                detection_information = detection_information + ', ' + str(colorstr('green', 'Region' + str(x + 1) + ': ')) + str(colorstr('blue', str(number_of_people[x]) + ' persons'))
            LOGGER.info(detection_information)
        if view_img:
            if vid_frame_count == 1:
                cv2.namedWindow('Region Counter Movable Powered By Anperlanch')
                cv2.setMouseCallback('Region Counter Movable Powered By Anperlanch', mouse_callback)
            cv2.imshow('Region Counter Movable Powered By Anperlanch', frame)
        if save_img:
            video_writer.write(frame)
        for region in counting_regions:
            region['counts'] = 0
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    workbook.save(r"runs\regional_counting.xlsx")
    for x in range(number_of_regions):
        LOGGER.info(f"{colorstr('yellow', 'Region' + str(x + 1) + ': ')}{colorstr('cyan', 'Maximum number of people:')} {colorstr('blue', str(record_table[x][1]) + ' persons')}, {colorstr('magenta', 'Time:')} {colorstr('red', record_table[x][0])}. {colorstr('cyan', 'Minimum number of people:')} {colorstr('blue', str(record_table[x][3]) + ' persons')}, {colorstr('magenta', 'Time:')} {colorstr('red', record_table[x][2])}")
    video_writer.release()
    videocapture.release()
    cv2.destroyAllWindows()


if __name__ == '__main__':
    video_path = 'audioAndVideoFiles/video.mp4'
    video_start_time = '2021-08-09 14:07:25'
    record_interval_time = 1
    x1, y1, height1, width1 = 110, 10, 600, 700
    x2, y2, height2, width2 = 820, 10, 600, 345
    x3, y3, height3, width3 = 105, 5, 610, 1070
    weights = r"C:\Users\15231\.conda\envs\yolo\Lib\site-packages\ultralytics\models\yolo\detect\yolov8n.pt"
    device = '0'
    view_img = True
    save_img = True
    exist_ok = True
    classes = 0
    line_thickness = 2
    track_thickness = 2
    region_thickness = 2
    main()

对摄像头拍到的实时视频进行区域计数

from collections import defaultdict
import cv2
import numpy as np
from shapely.geometry import Polygon
from shapely.geometry.point import Point
from ultralytics import YOLO
from ultralytics.utils.plotting import Annotator, colors
from datetime import datetime, timedelta
from ultralytics.utils import LOGGER, colorstr
import openpyxl
import re

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),
        }
    ]
    cv2.namedWindow('Region Counter Movable Powered By Anperlanch')
    cv2.setMouseCallback('Region Counter Movable Powered By Anperlanch', mouse_callback)
    i = 1
    model = YOLO(weights)
    model.to('cuda') if device == '0' else model.to('cpu')
    names = model.model.names
    videocapture = cv2.VideoCapture(0)
    video_writer = cv2.VideoWriter('runs/regional_counting.mp4', cv2.VideoWriter_fourcc(*'mp4v'), int(videocapture.get(5)), (int(videocapture.get(3)), int(videocapture.get(4))))
    workbook = openpyxl.Workbook()
    worksheet = workbook.active
    while videocapture.isOpened():
        success, frame = videocapture.read()
        if not success:
            break
        current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        if i == 1:
            printed_time = current_time
        cv2.putText(frame, current_time, (10, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        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
                track = track_history[track_id]
                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)
                for region in counting_regions:
                    if region['polygon'].contains(Point((bbox_center[0], bbox_center[1]))):
                        region['counts'] += 1
        number_of_people = []
        for region in counting_regions:
            number_of_people.append(region['counts'])
            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 current_time >= printed_time:
            year, month, day, hours, minutes, seconds = map(int, re.split('[-:\s]+', printed_time))
            printed_time = (datetime(year, month, day, hours, minutes, seconds) + timedelta(seconds=record_interval_time)).strftime('%Y-%m-%d %H:%M:%S')
            worksheet.append([current_time] + number_of_people)
            if i == 1:
                number_of_regions = len(counting_regions)
                record_table = [[0 for x in range(4)] for y in range(number_of_regions)]
                for x in range(number_of_regions):
                    for y in range(4):
                        if y == 0 or y == 2:
                            record_table[x][y] = current_time
                        else:
                            record_table[x][y] = number_of_people[x]
            for x in range(number_of_regions):
                if number_of_people[x] > record_table[x][1]:
                    record_table[x][1] = number_of_people[x]
                    record_table[x][0] = current_time
                if number_of_people[x] < record_table[x][3]:
                    record_table[x][3] = number_of_people[x]
                    record_table[x][2] = current_time
            i += 1
            detection_information = str(colorstr('magenta', "Time: ")) + str(colorstr('red', current_time))
            for x in range(number_of_regions):
                detection_information = detection_information + ', ' + str(colorstr('green', 'Region' + str(x + 1) + ': ')) + str(colorstr('blue', str(number_of_people[x]) + ' persons'))
            LOGGER.info(detection_information)
        cv2.imshow('Region Counter Movable Powered By Anperlanch', frame)
        video_writer.write(frame)
        for region in counting_regions:
            region['counts'] = 0
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    workbook.save(r"runs\regional_counting.xlsx")
    for x in range(number_of_regions):
        LOGGER.info(f"{colorstr('yellow', 'Region' + str(x + 1) + ': ')}{colorstr('cyan', 'Maximum number of people:')} {colorstr('blue', str(record_table[x][1]) + ' persons')}, {colorstr('magenta', 'Time:')} {colorstr('red', record_table[x][0])}. {colorstr('cyan', 'Minimum number of people:')} {colorstr('blue', str(record_table[x][3]) + ' persons')}, {colorstr('magenta', 'Time:')} {colorstr('red', record_table[x][2])}")
    video_writer.release()
    videocapture.release()
    cv2.destroyAllWindows()


if __name__ == '__main__':
    record_interval_time = 1
    x1, y1, height1, width1 = 10, 10, 460, 305
    x2, y2, height2, width2 = 325, 10, 460, 305
    x3, y3, height3, width3 = 5, 5, 470, 630
    weights = r"C:\Users\15231\.conda\envs\yolo\Lib\site-packages\ultralytics\models\yolo\detect\yolov8n.pt"
    device = '0'
    classes = 0
    line_thickness = 2
    track_thickness = 2
    region_thickness = 2
    main()

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值