ultralytics 8.2.17__py3-none-any.whl → 8.2.19__py3-none-any.whl

ultralytics/solutions/object_counter.py

@@ -15,55 +15,10 @@ from shapely.geometry import LineString, Point, Polygon
 class ObjectCounter:
     """A class to manage the counting of objects in a real-time video stream based on their tracks."""
 
-    def __init__(self):
-        """Initializes the Counter with default values for various tracking and counting parameters."""
-
-        # Mouse events
-        self.is_drawing = False
-        self.selected_point = None
-
-        # Region & Line Information
-        self.reg_pts = [(20, 400), (1260, 400)]
-        self.line_dist_thresh = 15
-        self.counting_region = None
-        self.region_color = (255, 0, 255)
-        self.region_thickness = 5
-
-        # Image and annotation Information
-        self.im0 = None
-        self.tf = None
-        self.view_img = False
-        self.view_in_counts = True
-        self.view_out_counts = True
-
-        self.names = None  # Classes names
-        self.annotator = None  # Annotator
-        self.window_name = "Ultralytics YOLOv8 Object Counter"
-
-        # Object counting Information
-        self.in_counts = 0
-        self.out_counts = 0
-        self.count_ids = []
-        self.class_wise_count = {}
-        self.count_txt_thickness = 0
-        self.count_txt_color = (255, 255, 255)
-        self.count_bg_color = (255, 255, 255)
-        self.cls_txtdisplay_gap = 50
-        self.fontsize = 0.6
-
-        # Tracks info
-        self.track_history = defaultdict(list)
-        self.track_thickness = 2
-        self.draw_tracks = False
-        self.track_color = None
-
-        # Check if environment support imshow
-        self.env_check = check_imshow(warn=True)
-
-    def set_args(
+    def __init__(
         self,
         classes_names,
-        reg_pts,
+        reg_pts=None,
         count_reg_color=(255, 0, 255),
         count_txt_color=(0, 0, 0),
         count_bg_color=(255, 255, 255),
@@ -79,66 +34,90 @@ class ObjectCounter:
         cls_txtdisplay_gap=50,
     ):
         """
-        Configures the Counter's image, bounding box line thickness, and counting region points.
+        Initializes the ObjectCounter with various tracking and counting parameters.
 
         Args:
+            classes_names (dict): Dictionary of class names.
+            reg_pts (list): List of points defining the counting region.
+            count_reg_color (tuple): RGB color of the counting region.
+            count_txt_color (tuple): RGB color of the count text.
+            count_bg_color (tuple): RGB color of the count text background.
             line_thickness (int): Line thickness for bounding boxes.
+            track_thickness (int): Thickness of the track lines.
             view_img (bool): Flag to control whether to display the video stream.
-            view_in_counts (bool): Flag to control whether to display the incounts on video stream.
-            view_out_counts (bool): Flag to control whether to display the outcounts on video stream.
-            reg_pts (list): Initial list of points defining the counting region.
-            classes_names (dict): Classes names
-            track_thickness (int): Track thickness
-            draw_tracks (Bool): draw tracks
-            count_txt_color (RGB color): count text color value
-            count_bg_color (RGB color): count highlighter line color
-            count_reg_color (RGB color): Color of object counting region
-            track_color (RGB color): color for tracks
-            region_thickness (int): Object counting Region thickness
-            line_dist_thresh (int): Euclidean Distance threshold for line counter
-            cls_txtdisplay_gap (int): Display gap between each class count
+            view_in_counts (bool): Flag to control whether to display the in counts on the video stream.
+            view_out_counts (bool): Flag to control whether to display the out counts on the video stream.
+            draw_tracks (bool): Flag to control whether to draw the object tracks.
+            track_color (tuple): RGB color of the tracks.
+            region_thickness (int): Thickness of the object counting region.
+            line_dist_thresh (int): Euclidean distance threshold for line counter.
+            cls_txtdisplay_gap (int): Display gap between each class count.
         """
+
+        # Mouse events
+        self.is_drawing = False
+        self.selected_point = None
+
+        # Region & Line Information
+        self.reg_pts = [(20, 400), (1260, 400)] if reg_pts is None else reg_pts
+        self.line_dist_thresh = line_dist_thresh
+        self.counting_region = None
+        self.region_color = count_reg_color
+        self.region_thickness = region_thickness
+
+        # Image and annotation Information
+        self.im0 = None
         self.tf = line_thickness
         self.view_img = view_img
         self.view_in_counts = view_in_counts
         self.view_out_counts = view_out_counts
+
+        self.names = classes_names  # Classes names
+        self.annotator = None  # Annotator
+        self.window_name = "Ultralytics YOLOv8 Object Counter"
+
+        # Object counting Information
+        self.in_counts = 0
+        self.out_counts = 0
+        self.count_ids = []
+        self.class_wise_count = {}
+        self.count_txt_thickness = 0
+        self.count_txt_color = count_txt_color
+        self.count_bg_color = count_bg_color
+        self.cls_txtdisplay_gap = cls_txtdisplay_gap
+        self.fontsize = 0.6
+
+        # Tracks info
+        self.track_history = defaultdict(list)
         self.track_thickness = track_thickness
         self.draw_tracks = draw_tracks
+        self.track_color = track_color
+
+        # Check if environment supports imshow
+        self.env_check = check_imshow(warn=True)
 
-        # Region and line selection
-        if len(reg_pts) == 2:
+        # Initialize counting region
+        if len(self.reg_pts) == 2:
             print("Line Counter Initiated.")
-            self.reg_pts = reg_pts
             self.counting_region = LineString(self.reg_pts)
-        elif len(reg_pts) >= 3:
+        elif len(self.reg_pts) >= 3:
             print("Polygon Counter Initiated.")
-            self.reg_pts = reg_pts
             self.counting_region = Polygon(self.reg_pts)
         else:
             print("Invalid Region points provided, region_points must be 2 for lines or >= 3 for polygons.")
             print("Using Line Counter Now")
             self.counting_region = LineString(self.reg_pts)
 
-        self.names = classes_names
-        self.track_color = track_color
-        self.count_txt_color = count_txt_color
-        self.count_bg_color = count_bg_color
-        self.region_color = count_reg_color
-        self.region_thickness = region_thickness
-        self.line_dist_thresh = line_dist_thresh
-        self.cls_txtdisplay_gap = cls_txtdisplay_gap
-
     def mouse_event_for_region(self, event, x, y, flags, params):
         """
-        This function is designed to move region with mouse events in a real-time video stream.
+        Handles mouse events for defining and moving the counting region in a real-time video stream.
 
         Args:
             event (int): The type of mouse event (e.g., cv2.EVENT_MOUSEMOVE, cv2.EVENT_LBUTTONDOWN, etc.).
             x (int): The x-coordinate of the mouse pointer.
             y (int): The y-coordinate of the mouse pointer.
-            flags (int): Any flags associated with the event (e.g., cv2.EVENT_FLAG_CTRLKEY,
-                cv2.EVENT_FLAG_SHIFTKEY, etc.).
-            params (dict): Additional parameters you may want to pass to the function.
+            flags (int): Any associated event flags (e.g., cv2.EVENT_FLAG_CTRLKEY,  cv2.EVENT_FLAG_SHIFTKEY, etc.).
+            params (dict): Additional parameters for the function.
         """
         if event == cv2.EVENT_LBUTTONDOWN:
             for i, point in enumerate(self.reg_pts):
@@ -240,11 +219,11 @@ class ObjectCounter:
                 else:
                     labels_dict[str.capitalize(key)] = f"IN {value['IN']} OUT {value['OUT']}"
 
-        if labels_dict is not None:
+        if labels_dict:
             self.annotator.display_analytics(self.im0, labels_dict, self.count_txt_color, self.count_bg_color, 10)
 
     def display_frames(self):
-        """Display frame."""
+        """Displays the current frame with annotations and regions in a window."""
         if self.env_check:
             cv2.namedWindow(self.window_name)
             if len(self.reg_pts) == 4:  # only add mouse event If user drawn region
@@ -271,4 +250,5 @@ class ObjectCounter:
 
 
 if __name__ == "__main__":
-    ObjectCounter()
+    classes_names = {0: "person", 1: "car"}  # example class names
+    ObjectCounter(classes_names)

ultralytics/solutions/parking_management.py

@@ -8,17 +8,22 @@ from PIL import Image, ImageTk
 from ultralytics.utils.checks import check_imshow, check_requirements
 from ultralytics.utils.plotting import Annotator
 
-check_requirements("tkinter")
-import tkinter as tk
-
 
 class ParkingPtsSelection:
     def __init__(self, master):
-        """Initializes the UI for selecting parking zone points in a tkinter window."""
+        """
+        Initializes the UI for selecting parking zone points in a tkinter window.
+
+        Args:
+            master (tk.Tk): The main tkinter window object.
+        """
+        check_requirements("tkinter")
+        import tkinter as tk
+
         self.master = master
         master.title("Ultralytics Parking Zones Points Selector")
 
-        # Resizable false
+        # Disable window resizing
         master.resizable(False, False)
 
         # Setup canvas for image display
@@ -36,7 +41,6 @@ class ParkingPtsSelection:
         self.image_path = None
         self.image = None
         self.canvas_image = None
-        self.canvas = None
         self.bounding_boxes = []
         self.current_box = []
         self.img_width = 0
@@ -101,7 +105,6 @@ class ParkingPtsSelection:
         Args:
             box (list): Bounding box data
         """
-
         for i in range(4):
             x1, y1 = box[i]
             x2, y2 = box[(i + 1) % 4]
@@ -151,6 +154,17 @@ class ParkingManagement:
         available_region_color=(0, 0, 255),
         margin=10,
     ):
+        """
+        Initializes the parking management system with a YOLOv8 model and visualization settings.
+
+        Args:
+            model_path (str): Path to the YOLOv8 model.
+            txt_color (tuple): RGB color tuple for text.
+            bg_color (tuple): RGB color tuple for background.
+            occupied_region_color (tuple): RGB color tuple for occupied regions.
+            available_region_color (tuple): RGB color tuple for available regions.
+            margin (int): Margin for text display.
+        """
         # Model path and initialization
         self.model_path = model_path
         self.model = self.load_model()
@@ -166,7 +180,7 @@ class ParkingManagement:
         self.available_region_color = available_region_color
 
         self.window_name = "Ultralytics YOLOv8 Parking Management System"
-        # Check if environment support imshow
+        # Check if environment supports imshow
         self.env_check = check_imshow(warn=True)
 
     def load_model(self):
@@ -184,7 +198,6 @@ class ParkingManagement:
         Args:
             json_file (str): file that have all parking slot points
         """
-
         with open(json_file, "r") as json_file:
             return json.load(json_file)
 

ultralytics/solutions/queue_management.py

@@ -13,49 +13,12 @@ from shapely.geometry import Point, Polygon
 
 
 class QueueManager:
-    """A class to manage the queue management in real-time video stream based on their tracks."""
+    """A class to manage the queue in a real-time video stream based on object tracks."""
 
-    def __init__(self):
-        """Initializes the queue manager with default values for various tracking and counting parameters."""
-
-        # Mouse events
-        self.is_drawing = False
-        self.selected_point = None
-
-        # Region & Line Information
-        self.reg_pts = [(20, 60), (20, 680), (1120, 680), (1120, 60)]
-        self.counting_region = None
-        self.region_color = (255, 0, 255)
-        self.region_thickness = 5
-
-        # Image and annotation Information
-        self.im0 = None
-        self.tf = None
-        self.view_img = False
-        self.view_queue_counts = True
-        self.fontsize = 0.6
-
-        self.names = None  # Classes names
-        self.annotator = None  # Annotator
-        self.window_name = "Ultralytics YOLOv8 Queue Manager"
-
-        # Object counting Information
-        self.counts = 0
-        self.count_txt_color = (255, 255, 255)
-
-        # Tracks info
-        self.track_history = defaultdict(list)
-        self.track_thickness = 2
-        self.draw_tracks = False
-        self.track_color = None
-
-        # Check if environment support imshow
-        self.env_check = check_imshow(warn=True)
-
-    def set_args(
+    def __init__(
         self,
         classes_names,
-        reg_pts,
+        reg_pts=None,
         line_thickness=2,
         track_thickness=2,
         view_img=False,
@@ -68,48 +31,65 @@ class QueueManager:
         fontsize=0.7,
     ):
         """
-        Configures the Counter's image, bounding box line thickness, and counting region points.
+        Initializes the QueueManager with specified parameters for tracking and counting objects.
 
         Args:
-            line_thickness (int): Line thickness for bounding boxes.
-            view_img (bool): Flag to control whether to display the video stream.
-            view_queue_counts (bool): Flag to control whether to display the counts on video stream.
-            reg_pts (list): Initial list of points defining the counting region.
-            classes_names (dict): Classes names
-            region_color (RGB color): Color of queue region
-            track_thickness (int): Track thickness
-            draw_tracks (Bool): draw tracks
-            count_txt_color (RGB color): count text color value
-            track_color (RGB color): color for tracks
-            region_thickness (int): Object counting Region thickness
-            fontsize (float): Text display font size
+            classes_names (dict): A dictionary mapping class IDs to class names.
+            reg_pts (list of tuples, optional): Points defining the counting region polygon. Defaults to a predefined
+                rectangle.
+            line_thickness (int, optional): Thickness of the annotation lines. Defaults to 2.
+            track_thickness (int, optional): Thickness of the track lines. Defaults to 2.
+            view_img (bool, optional): Whether to display the image frames. Defaults to False.
+            region_color (tuple, optional): Color of the counting region lines (BGR). Defaults to (255, 0, 255).
+            view_queue_counts (bool, optional): Whether to display the queue counts. Defaults to True.
+            draw_tracks (bool, optional): Whether to draw tracks of the objects. Defaults to False.
+            count_txt_color (tuple, optional): Color of the count text (BGR). Defaults to (255, 255, 255).
+            track_color (tuple, optional): Color of the tracks. If None, different colors will be used for different
+                tracks. Defaults to None.
+            region_thickness (int, optional): Thickness of the counting region lines. Defaults to 5.
+            fontsize (float, optional): Font size for the text annotations. Defaults to 0.7.
         """
+
+        # Mouse events state
+        self.is_drawing = False
+        self.selected_point = None
+
+        # Region & Line Information
+        self.reg_pts = reg_pts if reg_pts is not None else [(20, 60), (20, 680), (1120, 680), (1120, 60)]
+        self.counting_region = (
+            Polygon(self.reg_pts) if len(self.reg_pts) >= 3 else Polygon([(20, 60), (20, 680), (1120, 680), (1120, 60)])
+        )
+        self.region_color = region_color
+        self.region_thickness = region_thickness
+
+        # Image and annotation Information
+        self.im0 = None
         self.tf = line_thickness
         self.view_img = view_img
         self.view_queue_counts = view_queue_counts
-        self.track_thickness = track_thickness
-        self.draw_tracks = draw_tracks
-        self.region_color = region_color
+        self.fontsize = fontsize
 
-        if len(reg_pts) >= 3:
-            print("Queue region initiated...")
-            self.reg_pts = reg_pts
-            self.counting_region = Polygon(self.reg_pts)
-        else:
-            print("Invalid region points provided...")
-            print("Using default region now....")
-            self.counting_region = Polygon(self.reg_pts)
+        self.names = classes_names  # Class names
+        self.annotator = None  # Annotator
+        self.window_name = "Ultralytics YOLOv8 Queue Manager"
 
-        self.names = classes_names
-        self.track_color = track_color
+        # Object counting Information
+        self.counts = 0
         self.count_txt_color = count_txt_color
-        self.region_thickness = region_thickness
-        self.fontsize = fontsize
+
+        # Tracks info
+        self.track_history = defaultdict(list)
+        self.track_thickness = track_thickness
+        self.draw_tracks = draw_tracks
+        self.track_color = track_color
+
+        # Check if environment supports imshow
+        self.env_check = check_imshow(warn=True)
 
     def extract_and_process_tracks(self, tracks):
         """Extracts and processes tracks for queue management in a video stream."""
 
-        # Annotator Init and queue region drawing
+        # Initialize annotator and draw the queue region
         self.annotator = Annotator(self.im0, self.tf, self.names)
 
         if tracks[0].boxes.id is not None:
@@ -122,48 +102,48 @@ class QueueManager:
                 # Draw bounding box
                 self.annotator.box_label(box, label=f"{self.names[cls]}#{track_id}", color=colors(int(track_id), True))
 
-                # Draw Tracks
+                # Update track history
                 track_line = self.track_history[track_id]
                 track_line.append((float((box[0] + box[2]) / 2), float((box[1] + box[3]) / 2)))
                 if len(track_line) > 30:
                     track_line.pop(0)
 
-                # Draw track trails
+                # Draw track trails if enabled
                 if self.draw_tracks:
                     self.annotator.draw_centroid_and_tracks(
                         track_line,
-                        color=self.track_color if self.track_color else colors(int(track_id), True),
+                        color=self.track_color or colors(int(track_id), True),
                         track_thickness=self.track_thickness,
                     )
 
                 prev_position = self.track_history[track_id][-2] if len(self.track_history[track_id]) > 1 else None
 
+                # Check if the object is inside the counting region
                 if len(self.reg_pts) >= 3:
                     is_inside = self.counting_region.contains(Point(track_line[-1]))
                     if prev_position is not None and is_inside:
                         self.counts += 1
 
-        label = "Queue Counts : " + str(self.counts)
-
+        # Display queue counts
+        label = f"Queue Counts : {str(self.counts)}"
         if label is not None:
             self.annotator.queue_counts_display(
                 label,
                 points=self.reg_pts,
                 region_color=self.region_color,
                 txt_color=self.count_txt_color,
-                fontsize=self.fontsize,
             )
 
-        self.counts = 0
+        self.counts = 0  # Reset counts after displaying
         self.display_frames()
 
     def display_frames(self):
-        """Display frame."""
+        """Displays the current frame with annotations."""
         if self.env_check:
             self.annotator.draw_region(reg_pts=self.reg_pts, thickness=self.region_thickness, color=self.region_color)
             cv2.namedWindow(self.window_name)
             cv2.imshow(self.window_name, self.im0)
-            # Break Window
+            # Close window on 'q' key press
             if cv2.waitKey(1) & 0xFF == ord("q"):
                 return
 
@@ -175,13 +155,14 @@ class QueueManager:
             im0 (ndarray): Current frame from the video stream.
             tracks (list): List of tracks obtained from the object tracking process.
         """
-        self.im0 = im0  # store image
-        self.extract_and_process_tracks(tracks)  # draw region even if no objects
+        self.im0 = im0  # Store the current frame
+        self.extract_and_process_tracks(tracks)  # Extract and process tracks
 
         if self.view_img:
-            self.display_frames()
+            self.display_frames()  # Display the frame if enabled
         return self.im0
 
 
 if __name__ == "__main__":
-    QueueManager()
+    classes_names = {0: "person", 1: "car"}  # example class names
+    queue_manager = QueueManager(classes_names)