# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license from __future__ import annotations from typing import Any import numpy as np from ultralytics.solutions.solutions import BaseSolution, SolutionAnnotator, SolutionResults from ultralytics.utils.plotting import colors class RegionCounter(BaseSolution): """A class for real-time counting of objects within user-defined regions in a video stream. This class inherits from `BaseSolution` and provides functionality to define polygonal regions in a video frame, track objects, and count those objects that pass through each defined region. Useful for applications requiring counting in specified areas, such as monitoring zones or segmented sections. Attributes: region_template (dict): Template for creating new counting regions with default attributes including name, polygon coordinates, and display colors. counting_regions (list): List storing all defined regions, where each entry is based on `region_template` and includes specific region settings like name, coordinates, and color. region_counts (dict[str, int]): Dictionary storing the count of objects for each named region. Methods: add_region: Add a new counting region with specified attributes. process: Process video frames to count objects in each region. initialize_regions: Initialize zones to count the objects in each one. Zones could be multiple as well. Examples: Initialize a RegionCounter and add a counting region >>> counter = RegionCounter() >>> counter.add_region("Zone1", [(100, 100), (200, 100), (200, 200), (100, 200)], (255, 0, 0), (255, 255, 255)) >>> results = counter.process(frame) >>> print(f"Total tracks: {results.total_tracks}") """ def __init__(self, **kwargs: Any) -> None: """Initialize the RegionCounter for real-time object counting in user-defined regions.""" super().__init__(**kwargs) self.region_template = { "name": "Default Region", "polygon": None, "prepared_polygon": None, "counts": 0, "region_color": (255, 255, 255), "text_color": (0, 0, 0), } self.region_counts = {} self.counting_regions = [] self.initialize_regions() def add_region( self, name: str, polygon_points: list[tuple], region_color: tuple[int, int, int], text_color: tuple[int, int, int], ) -> dict[str, Any]: """Add 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[int, int, int]): BGR color for region visualization. text_color (tuple[int, int, int]): BGR color for the text within the region. Returns: (dict[str, Any]): Region information including name, polygon, and display colors. Raises: ValueError: If `polygon_points` contains fewer than 3 points. """ if len(polygon_points) < 3: raise ValueError( f"RegionCounter requires regions with at least 3 points to form a polygon, " f"but got {len(polygon_points)} for '{name}'." ) polygon = self.Polygon(polygon_points) region = self.region_template.copy() region.update( { "name": name, "polygon": polygon, "prepared_polygon": self.prep(polygon), "region_color": region_color, "text_color": text_color, } ) self.counting_regions.append(region) return region def initialize_regions(self): """Initialize counting regions from `self.region`, using a default rectangle if no region is provided. `self.region` may be a list of points (a single region named "Region#01") or a dict mapping region names to point lists. """ if self.region is None: self.initialize_region() if not isinstance(self.region, dict): # Ensure self.region is initialized and structured as a dictionary self.region = {"Region#01": self.region} for i, (name, pts) in enumerate(self.region.items()): self.add_region(name, pts, colors(i, True), (255, 255, 255)) def process(self, im0: np.ndarray) -> SolutionResults: """Process the input frame to detect and count objects within each defined region. Args: im0 (np.ndarray): Input image frame where objects and regions are annotated. Returns: (SolutionResults): Contains processed image `plot_im`, 'total_tracks' (int, total number of tracked objects), and 'region_counts' (dict, counts of objects per region). """ self.extract_tracks(im0) annotator = SolutionAnnotator(im0, line_width=self.line_width) self.region_counts = {region["name"]: 0 for region in self.counting_regions} for box, cls, track_id, conf in zip(self.boxes, self.clss, self.track_ids, self.confs): annotator.box_label(box, label=self.adjust_box_label(cls, conf, track_id), color=colors(track_id, True)) x0, y0, x1, y1 = self.get_enclosing_box(box) center = self.Point(((x0 + x1) / 2, (y0 + y1) / 2)) for region in self.counting_regions: if region["prepared_polygon"].contains(center): region["counts"] += 1 self.region_counts[region["name"]] = region["counts"] # Display region counts for region in self.counting_regions: poly = region["polygon"] pts = list(map(tuple, np.array(poly.exterior.coords, dtype=np.int32))) (x1, y1), (x2, y2) = [(int(poly.centroid.x), int(poly.centroid.y))] * 2 annotator.draw_region(pts, region["region_color"], self.line_width * 2) annotator.adaptive_label( [x1, y1, x2, y2], label=str(region["counts"]), color=region["region_color"], txt_color=region["text_color"], margin=self.line_width * 4, shape="rect", ) region["counts"] = 0 # Reset for next frame plot_im = annotator.result() self.display_output(plot_im) return SolutionResults(plot_im=plot_im, total_tracks=len(self.track_ids), region_counts=self.region_counts)