# Ultralytics YOLO 🚀, AGPL-3.0 license 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): """ Handles mouse events for region manipulation. Parameters: event (int): The mouse event type (e.g., cv2.EVENT_LBUTTONDOWN). x (int): The x-coordinate of the mouse pointer. y (int): The y-coordinate of the mouse pointer. flags (int): Additional flags passed by OpenCV. param: Additional parameters passed to the callback (not used in this function). Global Variables: current_region (dict): A dictionary representing the current selected region. Mouse Events: - LBUTTONDOWN: Initiates dragging for the region containing the clicked point. - MOUSEMOVE: Moves the selected region if dragging is active. - LBUTTONUP: Ends dragging for the selected region. Notes: - This function is intended to be used as a callback for OpenCV mouse events. - Requires the existence of the 'counting_regions' list and the 'Polygon' class. Example: >>> cv2.setMouseCallback(window_name, mouse_callback) """ 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="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu") parser.add_argument("--source", type=str, required=True, help="video file path") parser.add_argument("--view-img", action="store_true", help="show results") parser.add_argument("--save-img", action="store_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)