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

ultralytics/solutions/distance_calculation.py

@@ -9,94 +9,78 @@ from ultralytics.utils.plotting import Annotator, colors
 
 
 class DistanceCalculation:
-    """A class to calculate distance between two objects in real-time video stream based on their tracks."""
+    """A class to calculate distance between two objects in a real-time video stream based on their tracks."""
 
-    def __init__(self):
-        """Initializes the distance calculation class with default values for Visual, Image, track and distance
-        parameters.
+    def __init__(
+        self,
+        names,
+        pixels_per_meter=10,
+        view_img=False,
+        line_thickness=2,
+        line_color=(255, 255, 0),
+        centroid_color=(255, 0, 255),
+    ):
         """
+        Initializes the DistanceCalculation class with the given parameters.
 
-        # Visual & im0 information
+        Args:
+            names (dict): Dictionary mapping class indices to class names.
+            pixels_per_meter (int, optional): Conversion factor from pixels to meters. Defaults to 10.
+            view_img (bool, optional): Flag to indicate if the video stream should be displayed. Defaults to False.
+            line_thickness (int, optional): Thickness of the lines drawn on the image. Defaults to 2.
+            line_color (tuple, optional): Color of the lines drawn on the image (BGR format). Defaults to (255, 255, 0).
+            centroid_color (tuple, optional): Color of the centroids drawn (BGR format). Defaults to (255, 0, 255).
+        """
+        # Visual & image information
         self.im0 = None
         self.annotator = None
-        self.view_img = False
-        self.line_color = (255, 255, 0)
-        self.centroid_color = (255, 0, 255)
+        self.view_img = view_img
+        self.line_color = line_color
+        self.centroid_color = centroid_color
 
-        # Predict/track information
+        # Prediction & tracking information
         self.clss = None
-        self.names = None
+        self.names = names
         self.boxes = None
-        self.line_thickness = 2
+        self.line_thickness = line_thickness
         self.trk_ids = None
 
         # Distance calculation information
         self.centroids = []
-        self.pixel_per_meter = 10
+        self.pixel_per_meter = pixels_per_meter
 
-        # Mouse event
+        # Mouse event information
         self.left_mouse_count = 0
         self.selected_boxes = {}
 
-        # Check if environment support imshow
+        # Check if environment supports imshow
         self.env_check = check_imshow(warn=True)
 
-    def set_args(
-        self,
-        names,
-        pixels_per_meter=10,
-        view_img=False,
-        line_thickness=2,
-        line_color=(255, 255, 0),
-        centroid_color=(255, 0, 255),
-    ):
-        """
-        Configures the distance calculation and display parameters.
-
-        Args:
-            names (dict): object detection classes names
-            pixels_per_meter (int): Number of pixels in meter
-            view_img (bool): Flag indicating frame display
-            line_thickness (int): Line thickness for bounding boxes.
-            line_color (RGB): color of centroids line
-            centroid_color (RGB): colors of bbox centroids
-        """
-        self.names = names
-        self.pixel_per_meter = pixels_per_meter
-        self.view_img = view_img
-        self.line_thickness = line_thickness
-        self.line_color = line_color
-        self.centroid_color = centroid_color
-
     def mouse_event_for_distance(self, event, x, y, flags, param):
         """
-        This function is designed to move region with mouse events in a real-time video stream.
+        Handles mouse events to select regions 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.).
-            param (dict): Additional parameters you may want to pass to the function.
+            event (int): Type of mouse event (e.g., cv2.EVENT_MOUSEMOVE, cv2.EVENT_LBUTTONDOWN, etc.).
+            x (int): X-coordinate of the mouse pointer.
+            y (int): Y-coordinate of the mouse pointer.
+            flags (int): Flags associated with the event (e.g., cv2.EVENT_FLAG_CTRLKEY, cv2.EVENT_FLAG_SHIFTKEY, etc.).
+            param (dict): Additional parameters passed to the function.
         """
-        global selected_boxes
-        global left_mouse_count
         if event == cv2.EVENT_LBUTTONDOWN:
             self.left_mouse_count += 1
             if self.left_mouse_count <= 2:
                 for box, track_id in zip(self.boxes, self.trk_ids):
                     if box[0] < x < box[2] and box[1] < y < box[3] and track_id not in self.selected_boxes:
-                        self.selected_boxes[track_id] = []
                         self.selected_boxes[track_id] = box
 
-        if event == cv2.EVENT_RBUTTONDOWN:
+        elif event == cv2.EVENT_RBUTTONDOWN:
             self.selected_boxes = {}
             self.left_mouse_count = 0
 
     def extract_tracks(self, tracks):
         """
-        Extracts results from the provided data.
+        Extracts tracking results from the provided data.
 
         Args:
             tracks (list): List of tracks obtained from the object tracking process.
@@ -105,55 +89,65 @@ class DistanceCalculation:
         self.clss = tracks[0].boxes.cls.cpu().tolist()
         self.trk_ids = tracks[0].boxes.id.int().cpu().tolist()
 
-    def calculate_centroid(self, box):
+    @staticmethod
+    def calculate_centroid(box):
         """
-        Calculate the centroid of bounding box.
+        Calculates the centroid of a bounding box.
 
         Args:
-            box (list): Bounding box data
+            box (list): Bounding box coordinates [x1, y1, x2, y2].
+
+        Returns:
+            (tuple): Centroid coordinates (x, y).
         """
         return int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2)
 
     def calculate_distance(self, centroid1, centroid2):
         """
-        Calculate distance between two centroids.
+        Calculates the distance between two centroids.
 
         Args:
-            centroid1 (point): First bounding box data
-            centroid2 (point): Second bounding box data
+            centroid1 (tuple): Coordinates of the first centroid (x, y).
+            centroid2 (tuple): Coordinates of the second centroid (x, y).
+
+        Returns:
+            (tuple): Distance in meters and millimeters.
         """
         pixel_distance = math.sqrt((centroid1[0] - centroid2[0]) ** 2 + (centroid1[1] - centroid2[1]) ** 2)
-        return pixel_distance / self.pixel_per_meter, (pixel_distance / self.pixel_per_meter) * 1000
+        distance_m = pixel_distance / self.pixel_per_meter
+        distance_mm = distance_m * 1000
+        return distance_m, distance_mm
 
     def start_process(self, im0, tracks):
         """
-        Calculate distance between two bounding boxes based on tracking data.
+        Processes the video frame and calculates the distance between two bounding boxes.
 
         Args:
-            im0 (nd array): Image
+            im0 (ndarray): The image frame.
             tracks (list): List of tracks obtained from the object tracking process.
+
+        Returns:
+            (ndarray): The processed image frame.
         """
         self.im0 = im0
         if tracks[0].boxes.id is None:
             if self.view_img:
                 self.display_frames()
-            return
-        self.extract_tracks(tracks)
+            return im0
 
-        self.annotator = Annotator(self.im0, line_width=2)
+        self.extract_tracks(tracks)
+        self.annotator = Annotator(self.im0, line_width=self.line_thickness)
 
         for box, cls, track_id in zip(self.boxes, self.clss, self.trk_ids):
             self.annotator.box_label(box, color=colors(int(cls), True), label=self.names[int(cls)])
 
             if len(self.selected_boxes) == 2:
-                for trk_id, _ in self.selected_boxes.items():
+                for trk_id in self.selected_boxes.keys():
                     if trk_id == track_id:
                         self.selected_boxes[track_id] = box
 
         if len(self.selected_boxes) == 2:
-            for trk_id, box in self.selected_boxes.items():
-                centroid = self.calculate_centroid(self.selected_boxes[trk_id])
-                self.centroids.append(centroid)
+            self.centroids = [self.calculate_centroid(self.selected_boxes[trk_id]) for trk_id in self.selected_boxes]
 
             distance_m, distance_mm = self.calculate_distance(self.centroids[0], self.centroids[1])
             self.annotator.plot_distance_and_line(
@@ -168,7 +162,7 @@ class DistanceCalculation:
         return im0
 
     def display_frames(self):
-        """Display frame."""
+        """Displays the current frame with annotations."""
         cv2.namedWindow("Ultralytics Distance Estimation")
         cv2.setMouseCallback("Ultralytics Distance Estimation", self.mouse_event_for_distance)
         cv2.imshow("Ultralytics Distance Estimation", self.im0)
@@ -178,4 +172,5 @@ class DistanceCalculation:
 
 
 if __name__ == "__main__":
-    DistanceCalculation()
+    names = {0: "person", 1: "car"}  # example class names
+    distance_calculation = DistanceCalculation(names)

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)