Add region counter as ultralytics solution (#17439)
Co-authored-by: UltralyticsAssistant <web@ultralytics.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>pull/17515/head
Muhammad Rizwan Munawar
2 weeks ago
committed by
GitHub
parent
c8ff53185f
commit
1a5c35366e
6 changed files with 180 additions and 38 deletions
@ -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 |
||||||
Loading…
Reference in new issue