Spaces:
Running
Running
File size: 5,245 Bytes
1999a98 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
from ultralytics.solutions.solutions import BaseSolution
from ultralytics.utils import LOGGER
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):
if not isinstance(reg_pts, list) or not all(isinstance(pt, tuple) for pt in reg_pts):
LOGGER.warning(f"Invalid region points for {region_name}: {reg_pts}")
continue # Skip invalid entries
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
|