ultralytics 8.2.16__py3-none-any.whl → 8.2.18__py3-none-any.whl

ultralytics/solutions/heatmap.py

@@ -16,28 +16,48 @@ from shapely.geometry import LineString, Point, Polygon
 class Heatmap:
     """A class to draw heatmaps in real-time video stream based on their tracks."""
 
-    def __init__(self):
+    def __init__(
+        self,
+        classes_names,
+        imw=0,
+        imh=0,
+        colormap=cv2.COLORMAP_JET,
+        heatmap_alpha=0.5,
+        view_img=False,
+        view_in_counts=True,
+        view_out_counts=True,
+        count_reg_pts=None,
+        count_txt_color=(0, 0, 0),
+        count_bg_color=(255, 255, 255),
+        count_reg_color=(255, 0, 255),
+        region_thickness=5,
+        line_dist_thresh=15,
+        line_thickness=2,
+        decay_factor=0.99,
+        shape="circle",
+    ):
         """Initializes the heatmap class with default values for Visual, Image, track, count and heatmap parameters."""
 
         # Visual information
         self.annotator = None
-        self.view_img = False
-        self.shape = "circle"
+        self.view_img = view_img
+        self.shape = shape
 
-        self.names = None  # Classes names
+        self.initialized = False
+        self.names = classes_names  # Classes names
 
         # Image information
-        self.imw = None
-        self.imh = None
+        self.imw = imw
+        self.imh = imh
         self.im0 = None
-        self.tf = 2
-        self.view_in_counts = True
-        self.view_out_counts = True
+        self.tf = line_thickness
+        self.view_in_counts = view_in_counts
+        self.view_out_counts = view_out_counts
 
         # Heatmap colormap and heatmap np array
-        self.colormap = None
+        self.colormap = colormap
         self.heatmap = None
-        self.heatmap_alpha = 0.5
+        self.heatmap_alpha = heatmap_alpha
 
         # Predict/track information
         self.boxes = None
@@ -46,112 +66,48 @@ class Heatmap:
         self.track_history = defaultdict(list)
 
         # Region & Line Information
-        self.count_reg_pts = None
         self.counting_region = None
-        self.line_dist_thresh = 15
-        self.region_thickness = 5
-        self.region_color = (255, 0, 255)
+        self.line_dist_thresh = line_dist_thresh
+        self.region_thickness = region_thickness
+        self.region_color = count_reg_color
 
         # Object Counting Information
         self.in_counts = 0
         self.out_counts = 0
         self.count_ids = []
         self.class_wise_count = {}
-        self.count_txt_color = (0, 0, 0)
-        self.count_bg_color = (255, 255, 255)
+        self.count_txt_color = count_txt_color
+        self.count_bg_color = count_bg_color
         self.cls_txtdisplay_gap = 50
 
         # Decay factor
-        self.decay_factor = 0.99
+        self.decay_factor = decay_factor
 
-        # Check if environment support imshow
+        # Check if environment supports imshow
         self.env_check = check_imshow(warn=True)
 
-    def set_args(
-        self,
-        imw,
-        imh,
-        classes_names=None,
-        colormap=cv2.COLORMAP_JET,
-        heatmap_alpha=0.5,
-        view_img=False,
-        view_in_counts=True,
-        view_out_counts=True,
-        count_reg_pts=None,
-        count_txt_color=(0, 0, 0),
-        count_bg_color=(255, 255, 255),
-        count_reg_color=(255, 0, 255),
-        region_thickness=5,
-        line_dist_thresh=15,
-        line_thickness=2,
-        decay_factor=0.99,
-        shape="circle",
-    ):
-        """
-        Configures the heatmap colormap, width, height and display parameters.
-
-        Args:
-            colormap (cv2.COLORMAP): The colormap to be set.
-            imw (int): The width of the frame.
-            imh (int): The height of the frame.
-            classes_names (dict): Classes names
-            line_thickness (int): Line thickness for bounding boxes.
-            heatmap_alpha (float): alpha value for heatmap display
-            view_img (bool): Flag indicating frame display
-            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.
-            count_reg_pts (list): Object counting region points
-            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
-            region_thickness (int): Object counting Region thickness
-            line_dist_thresh (int): Euclidean Distance threshold for line counter
-            decay_factor (float): value for removing heatmap area after object passed
-            shape (str): Heatmap shape, rect or circle shape supported
-        """
-        self.tf = line_thickness
-        self.names = classes_names
-        self.imw = imw
-        self.imh = imh
-        self.heatmap_alpha = heatmap_alpha
-        self.view_img = view_img
-        self.view_in_counts = view_in_counts
-        self.view_out_counts = view_out_counts
-        self.colormap = colormap
-
         # Region and line selection
-        if count_reg_pts is not None:
-            if len(count_reg_pts) == 2:
+        self.count_reg_pts = count_reg_pts
+        print(self.count_reg_pts)
+        if self.count_reg_pts is not None:
+            if len(self.count_reg_pts) == 2:
                 print("Line Counter Initiated.")
-                self.count_reg_pts = count_reg_pts
                 self.counting_region = LineString(self.count_reg_pts)
-            elif len(count_reg_pts) >= 3:
+            elif len(self.count_reg_pts) >= 3:
                 print("Polygon Counter Initiated.")
-                self.count_reg_pts = count_reg_pts
                 self.counting_region = Polygon(self.count_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.count_reg_pts)
 
-        # Heatmap new frame
-        self.heatmap = np.zeros((int(self.imh), int(self.imw)), dtype=np.float32)
-
-        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.decay_factor = decay_factor
-        self.line_dist_thresh = line_dist_thresh
-        self.shape = shape
-
-        # shape of heatmap, if not selected
+        # Shape of heatmap, if not selected
         if self.shape not in {"circle", "rect"}:
             print("Unknown shape value provided, 'circle' & 'rect' supported")
             print("Using Circular shape now")
             self.shape = "circle"
 
-    def extract_results(self, tracks):
+    def extract_results(self, tracks, _intialized=False):
         """
         Extracts results from the provided data.
 
@@ -171,18 +127,20 @@ class Heatmap:
             tracks (list): List of tracks obtained from the object tracking process.
         """
         self.im0 = im0
-        if tracks[0].boxes.id is None:
-            self.heatmap = np.zeros((int(self.imh), int(self.imw)), dtype=np.float32)
-            if self.view_img and self.env_check:
-                self.display_frames()
-            return im0
+
+        # Initialize heatmap only once
+        if not self.initialized:
+            self.heatmap = np.zeros((int(self.im0.shape[0]), int(self.im0.shape[1])), dtype=np.float32)
+            self.initialized = True
+
         self.heatmap *= self.decay_factor  # decay factor
+
         self.extract_results(tracks)
         self.annotator = Annotator(self.im0, self.tf, None)
 
-        if self.count_reg_pts is not None:
+        if self.track_ids is not None:
             # Draw counting region
-            if self.view_in_counts or self.view_out_counts:
+            if self.count_reg_pts is not None:
                 self.annotator.draw_region(
                     reg_pts=self.count_reg_pts, color=self.region_color, thickness=self.region_thickness
                 )
@@ -214,25 +172,12 @@ class Heatmap:
 
                 prev_position = self.track_history[track_id][-2] if len(self.track_history[track_id]) > 1 else None
 
-                # Count objects in any polygon
-                if len(self.count_reg_pts) >= 3:
-                    is_inside = self.counting_region.contains(Point(track_line[-1]))
-
-                    if prev_position is not None and is_inside and track_id not in self.count_ids:
-                        self.count_ids.append(track_id)
-
-                        if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
-                            self.in_counts += 1
-                            self.class_wise_count[self.names[cls]]["IN"] += 1
-                        else:
-                            self.out_counts += 1
-                            self.class_wise_count[self.names[cls]]["OUT"] += 1
-
-                # Count objects using line
-                elif len(self.count_reg_pts) == 2:
-                    if prev_position is not None and track_id not in self.count_ids:
-                        distance = Point(track_line[-1]).distance(self.counting_region)
-                        if distance < self.line_dist_thresh and track_id not in self.count_ids:
+                if self.count_reg_pts is not None:
+                    # Count objects in any polygon
+                    if len(self.count_reg_pts) >= 3:
+                        is_inside = self.counting_region.contains(Point(track_line[-1]))
+
+                        if prev_position is not None and is_inside and track_id not in self.count_ids:
                             self.count_ids.append(track_id)
 
                             if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
@@ -242,6 +187,22 @@ class Heatmap:
                                 self.out_counts += 1
                                 self.class_wise_count[self.names[cls]]["OUT"] += 1
 
+                    # Count objects using line
+                    elif len(self.count_reg_pts) == 2:
+                        if prev_position is not None and track_id not in self.count_ids:
+                            distance = Point(track_line[-1]).distance(self.counting_region)
+                            if distance < self.line_dist_thresh and track_id not in self.count_ids:
+                                self.count_ids.append(track_id)
+
+                                if (box[0] - prev_position[0]) * (
+                                    self.counting_region.centroid.x - prev_position[0]
+                                ) > 0:
+                                    self.in_counts += 1
+                                    self.class_wise_count[self.names[cls]]["IN"] += 1
+                                else:
+                                    self.out_counts += 1
+                                    self.class_wise_count[self.names[cls]]["OUT"] += 1
+
         else:
             for box, cls in zip(self.boxes, self.clss):
                 if self.shape == "circle":
@@ -258,26 +219,26 @@ class Heatmap:
                 else:
                     self.heatmap[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])] += 2
 
+        if self.count_reg_pts is not None:
+            labels_dict = {}
+
+            for key, value in self.class_wise_count.items():
+                if value["IN"] != 0 or value["OUT"] != 0:
+                    if not self.view_in_counts and not self.view_out_counts:
+                        continue
+                    elif not self.view_in_counts:
+                        labels_dict[str.capitalize(key)] = f"OUT {value['OUT']}"
+                    elif not self.view_out_counts:
+                        labels_dict[str.capitalize(key)] = f"IN {value['IN']}"
+                    else:
+                        labels_dict[str.capitalize(key)] = f"IN {value['IN']} OUT {value['OUT']}"
+
+            if labels_dict is not None:
+                self.annotator.display_analytics(self.im0, labels_dict, self.count_txt_color, self.count_bg_color, 10)
+
         # Normalize, apply colormap to heatmap and combine with original image
         heatmap_normalized = cv2.normalize(self.heatmap, None, 0, 255, cv2.NORM_MINMAX)
         heatmap_colored = cv2.applyColorMap(heatmap_normalized.astype(np.uint8), self.colormap)
-
-        labels_dict = {}
-
-        for key, value in self.class_wise_count.items():
-            if value["IN"] != 0 or value["OUT"] != 0:
-                if not self.view_in_counts and not self.view_out_counts:
-                    continue
-                elif not self.view_in_counts:
-                    labels_dict[str.capitalize(key)] = f"OUT {value['OUT']}"
-                elif not self.view_out_counts:
-                    labels_dict[str.capitalize(key)] = f"IN {value['IN']}"
-                else:
-                    labels_dict[str.capitalize(key)] = f"IN {value['IN']} OUT {value['OUT']}"
-
-        if labels_dict is not None:
-            self.annotator.display_analytics(self.im0, labels_dict, self.count_txt_color, self.count_bg_color, 10)
-
         self.im0 = cv2.addWeighted(self.im0, 1 - self.heatmap_alpha, heatmap_colored, self.heatmap_alpha, 0)
 
         if self.env_check and self.view_img:
@@ -294,4 +255,5 @@ class Heatmap:
 
 
 if __name__ == "__main__":
-    Heatmap()
+    classes_names = {0: "person", 1: "car"}  # example class names
+    heatmap = Heatmap(classes_names)

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)