Add region counter as ultralytics solution (#17439)

Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>

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",
 )

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

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)