From 1a5c35366ef4577b00c35f9e8c5d5d0f05a61859 Mon Sep 17 00:00:00 2001 From: Muhammad Rizwan Munawar Date: Thu, 14 Nov 2024 00:21:54 +0500 Subject: [PATCH] Add region counter as ultralytics solution (#17439) Co-authored-by: UltralyticsAssistant Co-authored-by: Glenn Jocher --- docs/en/guides/region-counting.md | 85 +++++++------ docs/en/reference/solutions/region_counter.md | 16 +++ mkdocs.yml | 1 + ultralytics/solutions/__init__.py | 2 + ultralytics/solutions/region_counter.py | 112 ++++++++++++++++++ ultralytics/solutions/solutions.py | 2 + 6 files changed, 180 insertions(+), 38 deletions(-) create mode 100644 docs/en/reference/solutions/region_counter.md create mode 100644 ultralytics/solutions/region_counter.py diff --git a/docs/en/guides/region-counting.md b/docs/en/guides/region-counting.md index a27c2b4e53..c8363d68d3 100644 --- a/docs/en/guides/region-counting.md +++ b/docs/en/guides/region-counting.md @@ -34,56 +34,65 @@ keywords: object counting, regions, YOLOv8, computer vision, Ultralytics, effici | ![People Counting in Different Region using Ultralytics YOLOv8](https://github.com/ultralytics/docs/releases/download/0/people-counting-different-region-ultralytics-yolov8.avif) | ![Crowd Counting in Different Region using Ultralytics YOLOv8](https://github.com/ultralytics/docs/releases/download/0/crowd-counting-different-region-ultralytics-yolov8.avif) | | People Counting in Different Region using Ultralytics YOLOv8 | Crowd Counting in Different Region using Ultralytics YOLOv8 | -## Steps to Run +!!! example "Region Counting Example" -### Step 1: Install Required Libraries + === "Python" -Begin by cloning the Ultralytics repository, installing dependencies, and navigating to the local directory using the provided commands in Step 2. + ```python + import cv2 + from ultralytics import solutions -```bash -# Clone Ultralytics repo -git clone https://github.com/ultralytics/ultralytics + cap = cv2.VideoCapture("Path/to/video/file.mp4") + assert cap.isOpened(), "Error reading video file" + w, h, fps = (int(cap.get(x)) for x in (cv2.CAP_PROP_FRAME_WIDTH, cv2.CAP_PROP_FRAME_HEIGHT, cv2.CAP_PROP_FPS)) -# Navigate to the local directory -cd ultralytics/examples/YOLOv8-Region-Counter -``` + # Define region points + # region_points = [(20, 400), (1080, 404), (1080, 360), (20, 360)] # Pass region as list -### Step 2: Run Region Counting Using Ultralytics YOLOv8 + # pass region as dictionary + region_points = { + "region-01": [(50, 50), (250, 50), (250, 250), (50, 250)], + "region-02": [(640, 640), (780, 640), (780, 720), (640, 720)] + } -Execute the following basic commands for inference. + # Video writer + video_writer = cv2.VideoWriter("region_counting.avi", cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h)) -???+ tip "Region is Movable" + # Init Object Counter + region = solutions.RegionCounter( + show=True, + region=region_points, + model="yolo11n.pt", + ) - During video playback, you can interactively move the region within the video by clicking and dragging using the left mouse button. + # Process video + while cap.isOpened(): + success, im0 = cap.read() + if not success: + print("Video frame is empty or video processing has been successfully completed.") + break + im0 = region.count(im0) + video_writer.write(im0) -```bash -# Save results -python yolov8_region_counter.py --source "path/to/video.mp4" --save-img + cap.release() + video_writer.release() + cv2.destroyAllWindows() + ``` -# Run model on CPU -python yolov8_region_counter.py --source "path/to/video.mp4" --device cpu +!!! tip "Ultralytics Example Code" -# Change model file -python yolov8_region_counter.py --source "path/to/video.mp4" --weights "path/to/model.pt" + The Ultralytics region counting module is available in our [examples section](https://github.com/ultralytics/ultralytics/blob/main/examples/YOLOv8-Region-Counter/yolov8_region_counter.py). You can explore this example for code customization and modify it to suit your specific use case. -# Detect specific classes (e.g., first and third classes) -python yolov8_region_counter.py --source "path/to/video.mp4" --classes 0 2 +### Argument `RegionCounter` -# View results without saving -python yolov8_region_counter.py --source "path/to/video.mp4" --view-img -``` +Here's a table with the `RegionCounter` arguments: -### Optional Arguments - -| Name | Type | Default | Description | -| -------------------- | ------ | ------------ | --------------------------------------------------------------------------- | -| `--source` | `str` | `None` | Path to video file, for webcam 0 | -| `--line_thickness` | `int` | `2` | [Bounding Box](https://www.ultralytics.com/glossary/bounding-box) thickness | -| `--save-img` | `bool` | `False` | Save the predicted video/image | -| `--weights` | `str` | `yolov8n.pt` | Weights file path | -| `--classes` | `list` | `None` | Detect specific classes i.e. --classes 0 2 | -| `--region-thickness` | `int` | `2` | Region Box thickness | -| `--track-thickness` | `int` | `2` | Tracking line thickness | +| Name | Type | Default | Description | +| ------------ | ------ | -------------------------- | ---------------------------------------------------- | +| `model` | `str` | `None` | Path to Ultralytics YOLO Model File | +| `region` | `list` | `[(20, 400), (1260, 400)]` | List of points defining the counting region. | +| `line_width` | `int` | `2` | Line thickness for bounding boxes. | +| `show` | `bool` | `False` | Flag to control whether to display the video stream. | ## FAQ @@ -107,7 +116,7 @@ Follow these steps to run object counting in Ultralytics YOLOv8: python yolov8_region_counter.py --source "path/to/video.mp4" --save-img ``` -For more options, visit the [Run Region Counting](#steps-to-run) section. +For more options, visit the [Run Region Counting](https://github.com/ultralytics/ultralytics/blob/main/examples/YOLOv8-Region-Counter/readme.md) section. ### Why should I use Ultralytics YOLOv8 for object counting in regions? @@ -121,7 +130,7 @@ Explore deeper benefits in the [Advantages](#advantages-of-object-counting-in-re ### Can the defined regions be adjusted during video playback? -Yes, with Ultralytics YOLOv8, regions can be interactively moved during video playback. Simply click and drag with the left mouse button to reposition the region. This feature enhances flexibility for dynamic environments. Learn more in the tip section for [movable regions](#step-2-run-region-counting-using-ultralytics-yolov8). +Yes, with Ultralytics YOLOv8, regions can be interactively moved during video playback. Simply click and drag with the left mouse button to reposition the region. This feature enhances flexibility for dynamic environments. Learn more in the tip section for [movable regions](https://github.com/ultralytics/ultralytics/blob/33cdaa5782efb2bc2b5ede945771ba647882830d/examples/YOLOv8-Region-Counter/yolov8_region_counter.py#L39). ### What are some real-world applications of object counting in regions? diff --git a/docs/en/reference/solutions/region_counter.md b/docs/en/reference/solutions/region_counter.md new file mode 100644 index 0000000000..0f27adff28 --- /dev/null +++ b/docs/en/reference/solutions/region_counter.md @@ -0,0 +1,16 @@ +--- +description: Explore the Ultralytics Object Counter for real-time video streams. Learn about initializing parameters, tracking objects, and more. +keywords: Ultralytics, Object Counter, Real-time Tracking, Video Stream, Python, Object Detection +--- + +# Reference for `ultralytics/solutions/region_counter.py` + +!!! note + + This file is available at [https://github.com/ultralytics/ultralytics/blob/main/ultralytics/solutions/region_counter.py](https://github.com/ultralytics/ultralytics/blob/main/ultralytics/solutions/region_counter.py). If you spot a problem please help fix it by [contributing](https://docs.ultralytics.com/help/contributing/) a [Pull Request](https://github.com/ultralytics/ultralytics/edit/main/ultralytics/solutions/region_counter.py) 🛠️. Thank you 🙏! + +
+ +## ::: ultralytics.solutions.region_counter.RegionCounter + +

diff --git a/mkdocs.yml b/mkdocs.yml index 04d734430d..33bbbc3036 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -571,6 +571,7 @@ nav: - solutions: reference/solutions/solutions.md - speed_estimation: reference/solutions/speed_estimation.md - streamlit_inference: reference/solutions/streamlit_inference.md + - region_counter: reference/solutions/region_counter.md - trackers: - basetrack: reference/trackers/basetrack.md - bot_sort: reference/trackers/bot_sort.md diff --git a/ultralytics/solutions/__init__.py b/ultralytics/solutions/__init__.py index 4446c1826e..9de61edce2 100644 --- a/ultralytics/solutions/__init__.py +++ b/ultralytics/solutions/__init__.py @@ -7,6 +7,7 @@ from .heatmap import Heatmap from .object_counter import ObjectCounter from .parking_management import ParkingManagement, ParkingPtsSelection from .queue_management import QueueManager +from .region_counter import RegionCounter from .speed_estimation import SpeedEstimator from .streamlit_inference import inference @@ -21,4 +22,5 @@ __all__ = ( "SpeedEstimator", "Analytics", "inference", + "RegionCounter", ) diff --git a/ultralytics/solutions/region_counter.py b/ultralytics/solutions/region_counter.py new file mode 100644 index 0000000000..03575100d5 --- /dev/null +++ b/ultralytics/solutions/region_counter.py @@ -0,0 +1,112 @@ +# Ultralytics YOLO 🚀, AGPL-3.0 license + +from ultralytics.solutions.solutions import BaseSolution +from ultralytics.utils.plotting import Annotator, colors + + +class RegionCounter(BaseSolution): + """ + A class designed for real-time counting of objects within user-defined regions in a video stream. + + This class inherits from `BaseSolution` and offers functionalities to define polygonal regions in a video + frame, track objects, and count those objects that pass through each defined region. This makes it useful + for applications that require counting in specified areas, such as monitoring zones or segmented sections. + + Attributes: + region_template (dict): A template for creating new counting regions with default attributes including + the name, polygon coordinates, and display colors. + counting_regions (list): A list storing all defined regions, where each entry is based on `region_template` + and includes specific region settings like name, coordinates, and color. + + Methods: + add_region: Adds a new counting region with specified attributes, such as the region's name, polygon points, + region color, and text color. + count: Processes video frames to count objects in each region, drawing regions and displaying counts + on the frame. Handles object detection, region definition, and containment checks. + """ + + def __init__(self, **kwargs): + """Initializes the RegionCounter class for real-time counting in different regions of the video streams.""" + super().__init__(**kwargs) + self.region_template = { + "name": "Default Region", + "polygon": None, + "counts": 0, + "dragging": False, + "region_color": (255, 255, 255), + "text_color": (0, 0, 0), + } + self.counting_regions = [] + + def add_region(self, name, polygon_points, region_color, text_color): + """ + Adds a new region to the counting list based on the provided template with specific attributes. + + Args: + name (str): Name assigned to the new region. + polygon_points (list[tuple]): List of (x, y) coordinates defining the region's polygon. + region_color (tuple): BGR color for region visualization. + text_color (tuple): BGR color for the text within the region. + """ + region = self.region_template.copy() + region.update( + { + "name": name, + "polygon": self.Polygon(polygon_points), + "region_color": region_color, + "text_color": text_color, + } + ) + self.counting_regions.append(region) + + def count(self, im0): + """ + Processes the input frame to detect and count objects within each defined region. + + Args: + im0 (numpy.ndarray): Input image frame where objects and regions are annotated. + + Returns: + im0 (numpy.ndarray): Processed image frame with annotated counting information. + """ + self.annotator = Annotator(im0, line_width=self.line_width) + self.extract_tracks(im0) + + # Region initialization and conversion + if self.region is None: + self.initialize_region() + regions = {"Region#01": self.region} + else: + regions = self.region if isinstance(self.region, dict) else {"Region#01": self.region} + + # Draw regions and process counts for each defined area + for idx, (region_name, reg_pts) in enumerate(regions.items(), start=1): + color = colors(idx, True) + self.annotator.draw_region(reg_pts=reg_pts, color=color, thickness=self.line_width * 2) + self.add_region(region_name, reg_pts, color, self.annotator.get_txt_color()) + + # Prepare regions for containment check + for region in self.counting_regions: + region["prepared_polygon"] = self.prep(region["polygon"]) + + # Process bounding boxes and count objects within each region + for box, cls in zip(self.boxes, self.clss): + self.annotator.box_label(box, label=self.names[cls], color=colors(cls, True)) + bbox_center = ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2) + + for region in self.counting_regions: + if region["prepared_polygon"].contains(self.Point(bbox_center)): + region["counts"] += 1 + + # Display counts in each region + for region in self.counting_regions: + self.annotator.text_label( + region["polygon"].bounds, + label=str(region["counts"]), + color=region["region_color"], + txt_color=region["text_color"], + ) + region["counts"] = 0 # Reset count for next frame + + self.display_output(im0) + return im0 diff --git a/ultralytics/solutions/solutions.py b/ultralytics/solutions/solutions.py index ea94767033..ba3f1ec2cc 100644 --- a/ultralytics/solutions/solutions.py +++ b/ultralytics/solutions/solutions.py @@ -50,10 +50,12 @@ class BaseSolution: """ check_requirements("shapely>=2.0.0") from shapely.geometry import LineString, Point, Polygon + from shapely.prepared import prep self.LineString = LineString self.Polygon = Polygon self.Point = Point + self.prep = prep # Load config and update with args DEFAULT_SOL_DICT.update(kwargs)