update

ultralytics/solutions/__init__.py

@@ -0,0 +1 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license

ultralytics/solutions/ai_gym.py

@@ -0,0 +1,150 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
+import cv2
+
+from ultralytics.utils.checks import check_imshow
+from ultralytics.utils.plotting import Annotator
+
+
+class AIGym:
+    """A class to manage the gym steps of people in a real-time video stream based on their poses."""
+
+    def __init__(self):
+        """Initializes the AIGym with default values for Visual and Image parameters."""
+
+        # Image and line thickness
+        self.im0 = None
+        self.tf = None
+
+        # Keypoints and count information
+        self.keypoints = None
+        self.poseup_angle = None
+        self.posedown_angle = None
+        self.threshold = 0.001
+
+        # Store stage, count and angle information
+        self.angle = None
+        self.count = None
+        self.stage = None
+        self.pose_type = "pushup"
+        self.kpts_to_check = None
+
+        # Visual Information
+        self.view_img = False
+        self.annotator = None
+
+        # Check if environment support imshow
+        self.env_check = check_imshow(warn=True)
+
+    def set_args(
+        self,
+        kpts_to_check,
+        line_thickness=2,
+        view_img=False,
+        pose_up_angle=145.0,
+        pose_down_angle=90.0,
+        pose_type="pullup",
+    ):
+        """
+        Configures the AIGym line_thickness, save image and view image parameters.
+
+        Args:
+            kpts_to_check (list): 3 keypoints for counting
+            line_thickness (int): Line thickness for bounding boxes.
+            view_img (bool): display the im0
+            pose_up_angle (float): Angle to set pose position up
+            pose_down_angle (float): Angle to set pose position down
+            pose_type (str): "pushup", "pullup" or "abworkout"
+        """
+        self.kpts_to_check = kpts_to_check
+        self.tf = line_thickness
+        self.view_img = view_img
+        self.poseup_angle = pose_up_angle
+        self.posedown_angle = pose_down_angle
+        self.pose_type = pose_type
+
+    def start_counting(self, im0, results, frame_count):
+        """
+        Function used to count the gym steps.
+
+        Args:
+            im0 (ndarray): Current frame from the video stream.
+            results (list): Pose estimation data
+            frame_count (int): store current frame count
+        """
+        self.im0 = im0
+        if frame_count == 1:
+            self.count = [0] * len(results[0])
+            self.angle = [0] * len(results[0])
+            self.stage = ["-" for _ in results[0]]
+        self.keypoints = results[0].keypoints.data
+        self.annotator = Annotator(im0, line_width=2)
+
+        for ind, k in enumerate(reversed(self.keypoints)):
+            if self.pose_type in ["pushup", "pullup"]:
+                self.angle[ind] = self.annotator.estimate_pose_angle(
+                    k[int(self.kpts_to_check[0])].cpu(),
+                    k[int(self.kpts_to_check[1])].cpu(),
+                    k[int(self.kpts_to_check[2])].cpu(),
+                )
+                self.im0 = self.annotator.draw_specific_points(k, self.kpts_to_check, shape=(640, 640), radius=10)
+
+            if self.pose_type == "abworkout":
+                self.angle[ind] = self.annotator.estimate_pose_angle(
+                    k[int(self.kpts_to_check[0])].cpu(),
+                    k[int(self.kpts_to_check[1])].cpu(),
+                    k[int(self.kpts_to_check[2])].cpu(),
+                )
+                self.im0 = self.annotator.draw_specific_points(k, self.kpts_to_check, shape=(640, 640), radius=10)
+                if self.angle[ind] > self.poseup_angle:
+                    self.stage[ind] = "down"
+                if self.angle[ind] < self.posedown_angle and self.stage[ind] == "down":
+                    self.stage[ind] = "up"
+                    self.count[ind] += 1
+                self.annotator.plot_angle_and_count_and_stage(
+                    angle_text=self.angle[ind],
+                    count_text=self.count[ind],
+                    stage_text=self.stage[ind],
+                    center_kpt=k[int(self.kpts_to_check[1])],
+                    line_thickness=self.tf,
+                )
+
+            if self.pose_type == "pushup":
+                if self.angle[ind] > self.poseup_angle:
+                    self.stage[ind] = "up"
+                if self.angle[ind] < self.posedown_angle and self.stage[ind] == "up":
+                    self.stage[ind] = "down"
+                    self.count[ind] += 1
+                self.annotator.plot_angle_and_count_and_stage(
+                    angle_text=self.angle[ind],
+                    count_text=self.count[ind],
+                    stage_text=self.stage[ind],
+                    center_kpt=k[int(self.kpts_to_check[1])],
+                    line_thickness=self.tf,
+                )
+            if self.pose_type == "pullup":
+                if self.angle[ind] > self.poseup_angle:
+                    self.stage[ind] = "down"
+                if self.angle[ind] < self.posedown_angle and self.stage[ind] == "down":
+                    self.stage[ind] = "up"
+                    self.count[ind] += 1
+                self.annotator.plot_angle_and_count_and_stage(
+                    angle_text=self.angle[ind],
+                    count_text=self.count[ind],
+                    stage_text=self.stage[ind],
+                    center_kpt=k[int(self.kpts_to_check[1])],
+                    line_thickness=self.tf,
+                )
+
+            self.annotator.kpts(k, shape=(640, 640), radius=1, kpt_line=True)
+
+        if self.env_check and self.view_img:
+            cv2.imshow("Ultralytics YOLOv8 AI GYM", self.im0)
+            if cv2.waitKey(1) & 0xFF == ord("q"):
+                return
+
+        return self.im0
+
+
+if __name__ == "__main__":
+    AIGym()

ultralytics/solutions/distance_calculation.py

@@ -0,0 +1,181 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
+import math
+
+import cv2
+
+from ultralytics.utils.checks import check_imshow
+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."""
+
+    def __init__(self):
+        """Initializes the distance calculation class with default values for Visual, Image, track and distance
+        parameters.
+        """
+
+        # Visual & im0 information
+        self.im0 = None
+        self.annotator = None
+        self.view_img = False
+        self.line_color = (255, 255, 0)
+        self.centroid_color = (255, 0, 255)
+
+        # Predict/track information
+        self.clss = None
+        self.names = None
+        self.boxes = None
+        self.line_thickness = 2
+        self.trk_ids = None
+
+        # Distance calculation information
+        self.centroids = []
+        self.pixel_per_meter = 10
+
+        # Mouse event
+        self.left_mouse_count = 0
+        self.selected_boxes = {}
+
+        # Check if environment support 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.
+
+        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.
+        """
+        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:
+            self.selected_boxes = {}
+            self.left_mouse_count = 0
+
+    def extract_tracks(self, tracks):
+        """
+        Extracts results from the provided data.
+
+        Args:
+            tracks (list): List of tracks obtained from the object tracking process.
+        """
+        self.boxes = tracks[0].boxes.xyxy.cpu()
+        self.clss = tracks[0].boxes.cls.cpu().tolist()
+        self.trk_ids = tracks[0].boxes.id.int().cpu().tolist()
+
+    def calculate_centroid(self, box):
+        """
+        Calculate the centroid of bounding box.
+
+        Args:
+            box (list): Bounding box data
+        """
+        return int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2)
+
+    def calculate_distance(self, centroid1, centroid2):
+        """
+        Calculate distance between two centroids.
+
+        Args:
+            centroid1 (point): First bounding box data
+            centroid2 (point): Second bounding box data
+        """
+        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
+
+    def start_process(self, im0, tracks):
+        """
+        Calculate distance between two bounding boxes based on tracking data.
+
+        Args:
+            im0 (nd array): Image
+            tracks (list): List of tracks obtained from the object tracking process.
+        """
+        self.im0 = im0
+        if tracks[0].boxes.id is None:
+            if self.view_img:
+                self.display_frames()
+            return
+        self.extract_tracks(tracks)
+
+        self.annotator = Annotator(self.im0, line_width=2)
+
+        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():
+                    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)
+
+            distance_m, distance_mm = self.calculate_distance(self.centroids[0], self.centroids[1])
+            self.annotator.plot_distance_and_line(
+                distance_m, distance_mm, self.centroids, self.line_color, self.centroid_color
+            )
+
+        self.centroids = []
+
+        if self.view_img and self.env_check:
+            self.display_frames()
+
+        return im0
+
+    def display_frames(self):
+        """Display frame."""
+        cv2.namedWindow("Ultralytics Distance Estimation")
+        cv2.setMouseCallback("Ultralytics Distance Estimation", self.mouse_event_for_distance)
+        cv2.imshow("Ultralytics Distance Estimation", self.im0)
+
+        if cv2.waitKey(1) & 0xFF == ord("q"):
+            return
+
+
+if __name__ == "__main__":
+    DistanceCalculation()

ultralytics/solutions/heatmap.py

@@ -0,0 +1,281 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
+from collections import defaultdict
+
+import cv2
+import numpy as np
+
+from ultralytics.utils.checks import check_imshow, check_requirements
+from ultralytics.utils.plotting import Annotator
+
+check_requirements("shapely>=2.0.0")
+
+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):
+        """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"
+
+        # Image information
+        self.imw = None
+        self.imh = None
+        self.im0 = None
+        self.view_in_counts = True
+        self.view_out_counts = True
+
+        # Heatmap colormap and heatmap np array
+        self.colormap = None
+        self.heatmap = None
+        self.heatmap_alpha = 0.5
+
+        # Predict/track information
+        self.boxes = None
+        self.track_ids = None
+        self.clss = None
+        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)
+
+        # Object Counting Information
+        self.in_counts = 0
+        self.out_counts = 0
+        self.counting_list = []
+        self.count_txt_thickness = 0
+        self.count_txt_color = (0, 0, 0)
+        self.count_color = (255, 255, 255)
+
+        # Decay factor
+        self.decay_factor = 0.99
+
+        # Check if environment support imshow
+        self.env_check = check_imshow(warn=True)
+
+    def set_args(
+        self,
+        imw,
+        imh,
+        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_thickness=2,
+        count_txt_color=(0, 0, 0),
+        count_color=(255, 255, 255),
+        count_reg_color=(255, 0, 255),
+        region_thickness=5,
+        line_dist_thresh=15,
+        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.
+            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_thickness (int): Text thickness for object counting display
+            count_txt_color (RGB color): count text color value
+            count_color (RGB color): count text background color value
+            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.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:
+                print("Line Counter Initiated.")
+                self.count_reg_pts = count_reg_pts
+                self.counting_region = LineString(count_reg_pts)
+
+            elif len(count_reg_pts) == 4:
+                print("Region Counter Initiated.")
+                self.count_reg_pts = count_reg_pts
+                self.counting_region = Polygon(self.count_reg_pts)
+
+            else:
+                print("Region or line points Invalid, 2 or 4 points supported")
+                print("Using Line Counter Now")
+                self.counting_region = Polygon([(20, 400), (1260, 400)])  # dummy points
+
+        # Heatmap new frame
+        self.heatmap = np.zeros((int(self.imh), int(self.imw)), dtype=np.float32)
+
+        self.count_txt_thickness = count_txt_thickness
+        self.count_txt_color = count_txt_color
+        self.count_color = count_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
+        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):
+        """
+        Extracts results from the provided data.
+
+        Args:
+            tracks (list): List of tracks obtained from the object tracking process.
+        """
+        self.boxes = tracks[0].boxes.xyxy.cpu()
+        self.clss = tracks[0].boxes.cls.cpu().tolist()
+        self.track_ids = tracks[0].boxes.id.int().cpu().tolist()
+
+    def generate_heatmap(self, im0, tracks):
+        """
+        Generate heatmap based on tracking data.
+
+        Args:
+            im0 (nd array): Image
+            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
+        self.heatmap *= self.decay_factor  # decay factor
+        self.extract_results(tracks)
+        self.annotator = Annotator(self.im0, self.count_txt_thickness, None)
+
+        if self.count_reg_pts is not None:
+            # Draw counting region
+            if self.view_in_counts or self.view_out_counts:
+                self.annotator.draw_region(
+                    reg_pts=self.count_reg_pts, color=self.region_color, thickness=self.region_thickness
+                )
+
+            for box, cls, track_id in zip(self.boxes, self.clss, self.track_ids):
+                if self.shape == "circle":
+                    center = (int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2))
+                    radius = min(int(box[2]) - int(box[0]), int(box[3]) - int(box[1])) // 2
+
+                    y, x = np.ogrid[0 : self.heatmap.shape[0], 0 : self.heatmap.shape[1]]
+                    mask = (x - center[0]) ** 2 + (y - center[1]) ** 2 <= radius**2
+
+                    self.heatmap[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])] += (
+                        2 * mask[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])]
+                    )
+
+                else:
+                    self.heatmap[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])] += 2
+
+                # Store tracking hist
+                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)
+
+                # Count objects
+                if len(self.count_reg_pts) == 4:
+                    if self.counting_region.contains(Point(track_line[-1])) and track_id not in self.counting_list:
+                        self.counting_list.append(track_id)
+                        if box[0] < self.counting_region.centroid.x:
+                            self.out_counts += 1
+                        else:
+                            self.in_counts += 1
+
+                elif len(self.count_reg_pts) == 2:
+                    distance = Point(track_line[-1]).distance(self.counting_region)
+                    if distance < self.line_dist_thresh and track_id not in self.counting_list:
+                        self.counting_list.append(track_id)
+                        if box[0] < self.counting_region.centroid.x:
+                            self.out_counts += 1
+                        else:
+                            self.in_counts += 1
+        else:
+            for box, cls in zip(self.boxes, self.clss):
+                if self.shape == "circle":
+                    center = (int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2))
+                    radius = min(int(box[2]) - int(box[0]), int(box[3]) - int(box[1])) // 2
+
+                    y, x = np.ogrid[0 : self.heatmap.shape[0], 0 : self.heatmap.shape[1]]
+                    mask = (x - center[0]) ** 2 + (y - center[1]) ** 2 <= radius**2
+
+                    self.heatmap[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])] += (
+                        2 * mask[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])]
+                    )
+
+                else:
+                    self.heatmap[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])] += 2
+
+        # 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)
+
+        incount_label = f"In Count : {self.in_counts}"
+        outcount_label = f"OutCount : {self.out_counts}"
+
+        # Display counts based on user choice
+        counts_label = None
+        if not self.view_in_counts and not self.view_out_counts:
+            counts_label = None
+        elif not self.view_in_counts:
+            counts_label = outcount_label
+        elif not self.view_out_counts:
+            counts_label = incount_label
+        else:
+            counts_label = f"{incount_label} {outcount_label}"
+
+        if self.count_reg_pts is not None and counts_label is not None:
+            self.annotator.count_labels(
+                counts=counts_label,
+                count_txt_size=self.count_txt_thickness,
+                txt_color=self.count_txt_color,
+                color=self.count_color,
+            )
+
+        self.im0 = cv2.addWeighted(self.im0, 1 - self.heatmap_alpha, heatmap_colored, self.heatmap_alpha, 0)
+
+        if self.env_check and self.view_img:
+            self.display_frames()
+
+        return self.im0
+
+    def display_frames(self):
+        """Display frame."""
+        cv2.imshow("Ultralytics Heatmap", self.im0)
+
+        if cv2.waitKey(1) & 0xFF == ord("q"):
+            return
+
+
+if __name__ == "__main__":
+    Heatmap()

ultralytics/solutions/object_counter.py

@@ -0,0 +1,278 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
+from collections import defaultdict
+
+import cv2
+
+from ultralytics.utils.checks import check_imshow, check_requirements
+from ultralytics.utils.plotting import Annotator, colors
+
+check_requirements("shapely>=2.0.0")
+
+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.counting_dict = {}
+        self.count_txt_thickness = 0
+        self.count_txt_color = (0, 0, 0)
+        self.count_color = (255, 255, 255)
+
+        # Tracks info
+        self.track_history = defaultdict(list)
+        self.track_thickness = 2
+        self.draw_tracks = False
+        self.track_color = (0, 255, 0)
+
+        # Check if environment support imshow
+        self.env_check = check_imshow(warn=True)
+
+    def set_args(
+        self,
+        classes_names,
+        reg_pts,
+        count_reg_color=(255, 0, 255),
+        line_thickness=2,
+        track_thickness=2,
+        view_img=False,
+        view_in_counts=True,
+        view_out_counts=True,
+        draw_tracks=False,
+        count_txt_thickness=2,
+        count_txt_color=(0, 0, 0),
+        count_color=(255, 255, 255),
+        track_color=(0, 255, 0),
+        region_thickness=5,
+        line_dist_thresh=15,
+    ):
+        """
+        Configures the Counter's image, bounding box line thickness, and counting region points.
+
+        Args:
+            line_thickness (int): Line thickness for bounding boxes.
+            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_thickness (int): Text thickness for object counting display
+            count_txt_color (RGB color): count text color value
+            count_color (RGB color): count text background color value
+            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
+        """
+        self.tf = line_thickness
+        self.view_img = view_img
+        self.view_in_counts = view_in_counts
+        self.view_out_counts = view_out_counts
+        self.track_thickness = track_thickness
+        self.draw_tracks = draw_tracks
+
+        # Region and line selection
+        if len(reg_pts) == 2:
+            print("Line Counter Initiated.")
+            self.reg_pts = reg_pts
+            self.counting_region = LineString(self.reg_pts)
+        elif len(reg_pts) >= 3:
+            print("Region 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_thickness = count_txt_thickness
+        self.count_txt_color = count_txt_color
+        self.count_color = count_color
+        self.region_color = count_reg_color
+        self.region_thickness = region_thickness
+        self.line_dist_thresh = line_dist_thresh
+
+    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.
+
+        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.
+        """
+        if event == cv2.EVENT_LBUTTONDOWN:
+            for i, point in enumerate(self.reg_pts):
+                if (
+                    isinstance(point, (tuple, list))
+                    and len(point) >= 2
+                    and (abs(x - point[0]) < 10 and abs(y - point[1]) < 10)
+                ):
+                    self.selected_point = i
+                    self.is_drawing = True
+                    break
+
+        elif event == cv2.EVENT_MOUSEMOVE:
+            if self.is_drawing and self.selected_point is not None:
+                self.reg_pts[self.selected_point] = (x, y)
+                self.counting_region = Polygon(self.reg_pts)
+
+        elif event == cv2.EVENT_LBUTTONUP:
+            self.is_drawing = False
+            self.selected_point = None
+
+    def extract_and_process_tracks(self, tracks):
+        """Extracts and processes tracks for object counting in a video stream."""
+
+        # Annotator Init and region drawing
+        self.annotator = Annotator(self.im0, self.tf, self.names)
+
+        if tracks[0].boxes.id is not None:
+            boxes = tracks[0].boxes.xyxy.cpu()
+            clss = tracks[0].boxes.cls.cpu().tolist()
+            track_ids = tracks[0].boxes.id.int().cpu().tolist()
+
+            # Extract tracks
+            for box, track_id, cls in zip(boxes, track_ids, clss):
+                # Draw bounding box
+                self.annotator.box_label(box, label=f"{track_id}:{self.names[cls]}", color=colors(int(track_id), True))
+
+                # Draw Tracks
+                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
+                if self.draw_tracks:
+                    self.annotator.draw_centroid_and_tracks(
+                        track_line, color=self.track_color, track_thickness=self.track_thickness
+                    )
+
+                prev_position = self.track_history[track_id][-2] if len(self.track_history[track_id]) > 1 else None
+                centroid = Point((box[:2] + box[2:]) / 2)
+
+                # Count objects
+                if len(self.reg_pts) >= 3:  # any polygon
+                    is_inside = self.counting_region.contains(centroid)
+                    current_position = "in" if is_inside else "out"
+
+                    if prev_position is not None:
+                        if self.counting_dict[track_id] != current_position and is_inside:
+                            self.in_counts += 1
+                            self.counting_dict[track_id] = "in"
+                        elif self.counting_dict[track_id] != current_position and not is_inside:
+                            self.out_counts += 1
+                            self.counting_dict[track_id] = "out"
+                        else:
+                            self.counting_dict[track_id] = current_position
+
+                    else:
+                        self.counting_dict[track_id] = current_position
+
+                elif len(self.reg_pts) == 2:
+                    if prev_position is not None:
+                        is_inside = (box[0] - prev_position[0]) * (
+                            self.counting_region.centroid.x - prev_position[0]
+                        ) > 0
+                        current_position = "in" if is_inside else "out"
+
+                        if self.counting_dict[track_id] != current_position and is_inside:
+                            self.in_counts += 1
+                            self.counting_dict[track_id] = "in"
+                        elif self.counting_dict[track_id] != current_position and not is_inside:
+                            self.out_counts += 1
+                            self.counting_dict[track_id] = "out"
+                        else:
+                            self.counting_dict[track_id] = current_position
+                    else:
+                        self.counting_dict[track_id] = None
+
+        incount_label = f"In Count : {self.in_counts}"
+        outcount_label = f"OutCount : {self.out_counts}"
+
+        # Display counts based on user choice
+        counts_label = None
+        if not self.view_in_counts and not self.view_out_counts:
+            counts_label = None
+        elif not self.view_in_counts:
+            counts_label = outcount_label
+        elif not self.view_out_counts:
+            counts_label = incount_label
+        else:
+            counts_label = f"{incount_label} {outcount_label}"
+
+        if counts_label is not None:
+            self.annotator.count_labels(
+                counts=counts_label,
+                count_txt_size=self.count_txt_thickness,
+                txt_color=self.count_txt_color,
+                color=self.count_color,
+            )
+
+    def display_frames(self):
+        """Display frame."""
+        if self.env_check:
+            self.annotator.draw_region(reg_pts=self.reg_pts, color=self.region_color, thickness=self.region_thickness)
+            cv2.namedWindow(self.window_name)
+            if len(self.reg_pts) == 4:  # only add mouse event If user drawn region
+                cv2.setMouseCallback(self.window_name, self.mouse_event_for_region, {"region_points": self.reg_pts})
+            cv2.imshow(self.window_name, self.im0)
+            # Break Window
+            if cv2.waitKey(1) & 0xFF == ord("q"):
+                return
+
+    def start_counting(self, im0, tracks):
+        """
+        Main function to start the object counting process.
+
+        Args:
+            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
+
+        if self.view_img:
+            self.display_frames()
+        return self.im0
+
+
+if __name__ == "__main__":
+    ObjectCounter()

ultralytics/solutions/speed_estimation.py

@@ -0,0 +1,198 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
+from collections import defaultdict
+from time import time
+
+import cv2
+import numpy as np
+
+from ultralytics.utils.checks import check_imshow
+from ultralytics.utils.plotting import Annotator, colors
+
+
+class SpeedEstimator:
+    """A class to estimation speed of objects in real-time video stream based on their tracks."""
+
+    def __init__(self):
+        """Initializes the speed-estimator class with default values for Visual, Image, track and speed parameters."""
+
+        # Visual & im0 information
+        self.im0 = None
+        self.annotator = None
+        self.view_img = False
+
+        # Region information
+        self.reg_pts = [(20, 400), (1260, 400)]
+        self.region_thickness = 3
+
+        # Predict/track information
+        self.clss = None
+        self.names = None
+        self.boxes = None
+        self.trk_ids = None
+        self.trk_pts = None
+        self.line_thickness = 2
+        self.trk_history = defaultdict(list)
+
+        # Speed estimator information
+        self.current_time = 0
+        self.dist_data = {}
+        self.trk_idslist = []
+        self.spdl_dist_thresh = 10
+        self.trk_previous_times = {}
+        self.trk_previous_points = {}
+
+        # Check if environment support imshow
+        self.env_check = check_imshow(warn=True)
+
+    def set_args(
+        self,
+        reg_pts,
+        names,
+        view_img=False,
+        line_thickness=2,
+        region_thickness=5,
+        spdl_dist_thresh=10,
+    ):
+        """
+        Configures the speed estimation and display parameters.
+
+        Args:
+            reg_pts (list): Initial list of points defining the speed calculation region.
+            names (dict): object detection classes names
+            view_img (bool): Flag indicating frame display
+            line_thickness (int): Line thickness for bounding boxes.
+            region_thickness (int): Speed estimation region thickness
+            spdl_dist_thresh (int): Euclidean distance threshold for speed line
+        """
+        if reg_pts is None:
+            print("Region points not provided, using default values")
+        else:
+            self.reg_pts = reg_pts
+        self.names = names
+        self.view_img = view_img
+        self.line_thickness = line_thickness
+        self.region_thickness = region_thickness
+        self.spdl_dist_thresh = spdl_dist_thresh
+
+    def extract_tracks(self, tracks):
+        """
+        Extracts results from the provided data.
+
+        Args:
+            tracks (list): List of tracks obtained from the object tracking process.
+        """
+        self.boxes = tracks[0].boxes.xyxy.cpu()
+        self.clss = tracks[0].boxes.cls.cpu().tolist()
+        self.trk_ids = tracks[0].boxes.id.int().cpu().tolist()
+
+    def store_track_info(self, track_id, box):
+        """
+        Store track data.
+
+        Args:
+            track_id (int): object track id.
+            box (list): object bounding box data
+        """
+        track = self.trk_history[track_id]
+        bbox_center = (float((box[0] + box[2]) / 2), float((box[1] + box[3]) / 2))
+        track.append(bbox_center)
+
+        if len(track) > 30:
+            track.pop(0)
+
+        self.trk_pts = np.hstack(track).astype(np.int32).reshape((-1, 1, 2))
+        return track
+
+    def plot_box_and_track(self, track_id, box, cls, track):
+        """
+        Plot track and bounding box.
+
+        Args:
+            track_id (int): object track id.
+            box (list): object bounding box data
+            cls (str): object class name
+            track (list): tracking history for tracks path drawing
+        """
+        speed_label = f"{int(self.dist_data[track_id])}km/ph" if track_id in self.dist_data else self.names[int(cls)]
+        bbox_color = colors(int(track_id)) if track_id in self.dist_data else (255, 0, 255)
+
+        self.annotator.box_label(box, speed_label, bbox_color)
+
+        cv2.polylines(self.im0, [self.trk_pts], isClosed=False, color=(0, 255, 0), thickness=1)
+        cv2.circle(self.im0, (int(track[-1][0]), int(track[-1][1])), 5, bbox_color, -1)
+
+    def calculate_speed(self, trk_id, track):
+        """
+        Calculation of object speed.
+
+        Args:
+            trk_id (int): object track id.
+            track (list): tracking history for tracks path drawing
+        """
+
+        if not self.reg_pts[0][0] < track[-1][0] < self.reg_pts[1][0]:
+            return
+        if self.reg_pts[1][1] - self.spdl_dist_thresh < track[-1][1] < self.reg_pts[1][1] + self.spdl_dist_thresh:
+            direction = "known"
+
+        elif self.reg_pts[0][1] - self.spdl_dist_thresh < track[-1][1] < self.reg_pts[0][1] + self.spdl_dist_thresh:
+            direction = "known"
+
+        else:
+            direction = "unknown"
+
+        if self.trk_previous_times[trk_id] != 0 and direction != "unknown" and trk_id not in self.trk_idslist:
+            self.trk_idslist.append(trk_id)
+
+            time_difference = time() - self.trk_previous_times[trk_id]
+            if time_difference > 0:
+                dist_difference = np.abs(track[-1][1] - self.trk_previous_points[trk_id][1])
+                speed = dist_difference / time_difference
+                self.dist_data[trk_id] = speed
+
+        self.trk_previous_times[trk_id] = time()
+        self.trk_previous_points[trk_id] = track[-1]
+
+    def estimate_speed(self, im0, tracks, region_color=(255, 0, 0)):
+        """
+        Calculate object based on tracking data.
+
+        Args:
+            im0 (nd array): Image
+            tracks (list): List of tracks obtained from the object tracking process.
+            region_color (tuple): Color to use when drawing regions.
+        """
+        self.im0 = im0
+        if tracks[0].boxes.id is None:
+            if self.view_img and self.env_check:
+                self.display_frames()
+            return im0
+        self.extract_tracks(tracks)
+
+        self.annotator = Annotator(self.im0, line_width=2)
+        self.annotator.draw_region(reg_pts=self.reg_pts, color=region_color, thickness=self.region_thickness)
+
+        for box, trk_id, cls in zip(self.boxes, self.trk_ids, self.clss):
+            track = self.store_track_info(trk_id, box)
+
+            if trk_id not in self.trk_previous_times:
+                self.trk_previous_times[trk_id] = 0
+
+            self.plot_box_and_track(trk_id, box, cls, track)
+            self.calculate_speed(trk_id, track)
+
+        if self.view_img and self.env_check:
+            self.display_frames()
+
+        return im0
+
+    def display_frames(self):
+        """Display frame."""
+        cv2.imshow("Ultralytics Speed Estimation", self.im0)
+        if cv2.waitKey(1) & 0xFF == ord("q"):
+            return
+
+
+if __name__ == "__main__":
+    SpeedEstimator()