first commit

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,297 @@
+# 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"
+
+        self.names = None  # Classes names
+
+        # Image information
+        self.imw = None
+        self.imh = None
+        self.im0 = None
+        self.tf = 2
+        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.count_ids = []
+        self.class_wise_count = {}
+        self.count_txt_color = (0, 0, 0)
+        self.count_bg_color = (255, 255, 255)
+        self.cls_txtdisplay_gap = 50
+
+        # 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,
+        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:
+                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:
+                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
+        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.tf, 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):
+                # Store class info
+                if self.names[cls] not in self.class_wise_count:
+                    self.class_wise_count[self.names[cls]] = {"IN": 0, "OUT": 0}
+
+                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)
+
+                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:
+                            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":
+                    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)
+
+        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:
+            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,274 @@
+# 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.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(
+        self,
+        classes_names,
+        reg_pts,
+        count_reg_color=(255, 0, 255),
+        count_txt_color=(0, 0, 0),
+        count_bg_color=(255, 255, 255),
+        line_thickness=2,
+        track_thickness=2,
+        view_img=False,
+        view_in_counts=True,
+        view_out_counts=True,
+        draw_tracks=False,
+        track_color=None,
+        region_thickness=5,
+        line_dist_thresh=15,
+        cls_txtdisplay_gap=50,
+    ):
+        """
+        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_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
+        """
+        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("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.
+
+        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)
+
+        # Draw region or line
+        self.annotator.draw_region(reg_pts=self.reg_pts, color=self.region_color, thickness=self.region_thickness)
+
+        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"{self.names[cls]}#{track_id}", color=colors(int(track_id), True))
+
+                # Store class info
+                if self.names[cls] not in self.class_wise_count:
+                    self.class_wise_count[self.names[cls]] = {"IN": 0, "OUT": 0}
+
+                # 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 if self.track_color else 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
+
+                # Count objects in any polygon
+                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 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.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
+
+        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)
+
+    def display_frames(self):
+        """Display frame."""
+        if self.env_check:
+            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/parking_management.py

@@ -0,0 +1,249 @@
+import json
+from tkinter import filedialog, messagebox
+
+import cv2
+import numpy as np
+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."""
+        self.master = master
+        master.title("Ultralytics Parking Zones Points Selector")
+
+        # Resizable false
+        master.resizable(False, False)
+
+        # Setup canvas for image display
+        self.canvas = tk.Canvas(master, bg="white")
+
+        # Setup buttons
+        button_frame = tk.Frame(master)
+        button_frame.pack(side=tk.TOP)
+
+        tk.Button(button_frame, text="Upload Image", command=self.upload_image).grid(row=0, column=0)
+        tk.Button(button_frame, text="Remove Last BBox", command=self.remove_last_bounding_box).grid(row=0, column=1)
+        tk.Button(button_frame, text="Save", command=self.save_to_json).grid(row=0, column=2)
+
+        # Initialize properties
+        self.image_path = None
+        self.image = None
+        self.canvas_image = None
+        self.canvas = None
+        self.bounding_boxes = []
+        self.current_box = []
+        self.img_width = 0
+        self.img_height = 0
+
+        # Constants
+        self.canvas_max_width = 1280
+        self.canvas_max_height = 720
+
+    def upload_image(self):
+        """Upload an image and resize it to fit canvas."""
+        self.image_path = filedialog.askopenfilename(filetypes=[("Image Files", "*.png;*.jpg;*.jpeg")])
+        if not self.image_path:
+            return
+
+        self.image = Image.open(self.image_path)
+        self.img_width, self.img_height = self.image.size
+
+        # Calculate the aspect ratio and resize image
+        aspect_ratio = self.img_width / self.img_height
+        if aspect_ratio > 1:
+            # Landscape orientation
+            canvas_width = min(self.canvas_max_width, self.img_width)
+            canvas_height = int(canvas_width / aspect_ratio)
+        else:
+            # Portrait orientation
+            canvas_height = min(self.canvas_max_height, self.img_height)
+            canvas_width = int(canvas_height * aspect_ratio)
+
+        # Check if canvas is already initialized
+        if self.canvas:
+            self.canvas.destroy()  # Destroy previous canvas
+
+        self.canvas = tk.Canvas(self.master, bg="white", width=canvas_width, height=canvas_height)
+        resized_image = self.image.resize((canvas_width, canvas_height), Image.LANCZOS)
+        self.canvas_image = ImageTk.PhotoImage(resized_image)
+        self.canvas.create_image(0, 0, anchor=tk.NW, image=self.canvas_image)
+
+        self.canvas.pack(side=tk.BOTTOM)
+        self.canvas.bind("<Button-1>", self.on_canvas_click)
+
+        # Reset bounding boxes and current box
+        self.bounding_boxes = []
+        self.current_box = []
+
+    def on_canvas_click(self, event):
+        """Handle mouse clicks on canvas to create points for bounding boxes."""
+        self.current_box.append((event.x, event.y))
+        x0, y0 = event.x - 3, event.y - 3
+        x1, y1 = event.x + 3, event.y + 3
+        self.canvas.create_oval(x0, y0, x1, y1, fill="red")
+
+        if len(self.current_box) == 4:
+            self.bounding_boxes.append(self.current_box)
+            self.draw_bounding_box(self.current_box)
+            self.current_box = []
+
+    def draw_bounding_box(self, box):
+        """
+        Draw bounding box on canvas.
+
+        Args:
+            box (list): Bounding box data
+        """
+
+        for i in range(4):
+            x1, y1 = box[i]
+            x2, y2 = box[(i + 1) % 4]
+            self.canvas.create_line(x1, y1, x2, y2, fill="blue", width=2)
+
+    def remove_last_bounding_box(self):
+        """Remove the last drawn bounding box from canvas."""
+        if self.bounding_boxes:
+            self.bounding_boxes.pop()  # Remove the last bounding box
+            self.canvas.delete("all")  # Clear the canvas
+            self.canvas.create_image(0, 0, anchor=tk.NW, image=self.canvas_image)  # Redraw the image
+
+            # Redraw all bounding boxes
+            for box in self.bounding_boxes:
+                self.draw_bounding_box(box)
+
+            messagebox.showinfo("Success", "Last bounding box removed.")
+        else:
+            messagebox.showwarning("Warning", "No bounding boxes to remove.")
+
+    def save_to_json(self):
+        """Saves rescaled bounding boxes to 'bounding_boxes.json' based on image-to-canvas size ratio."""
+        canvas_width, canvas_height = self.canvas.winfo_width(), self.canvas.winfo_height()
+        width_scaling_factor = self.img_width / canvas_width
+        height_scaling_factor = self.img_height / canvas_height
+        bounding_boxes_data = []
+        for box in self.bounding_boxes:
+            rescaled_box = []
+            for x, y in box:
+                rescaled_x = int(x * width_scaling_factor)
+                rescaled_y = int(y * height_scaling_factor)
+                rescaled_box.append((rescaled_x, rescaled_y))
+            bounding_boxes_data.append({"points": rescaled_box})
+        with open("bounding_boxes.json", "w") as json_file:
+            json.dump(bounding_boxes_data, json_file, indent=4)
+
+        messagebox.showinfo("Success", "Bounding boxes saved to bounding_boxes.json")
+
+
+class ParkingManagement:
+    def __init__(
+        self,
+        model_path,
+        txt_color=(0, 0, 0),
+        bg_color=(255, 255, 255),
+        occupied_region_color=(0, 255, 0),
+        available_region_color=(0, 0, 255),
+        margin=10,
+    ):
+        # Model path and initialization
+        self.model_path = model_path
+        self.model = self.load_model()
+
+        # Labels dictionary
+        self.labels_dict = {"Occupancy": 0, "Available": 0}
+
+        # Visualization details
+        self.margin = margin
+        self.bg_color = bg_color
+        self.txt_color = txt_color
+        self.occupied_region_color = occupied_region_color
+        self.available_region_color = available_region_color
+
+        self.window_name = "Ultralytics YOLOv8 Parking Management System"
+        # Check if environment support imshow
+        self.env_check = check_imshow(warn=True)
+
+    def load_model(self):
+        """Load the Ultralytics YOLOv8 model for inference and analytics."""
+        from ultralytics import YOLO
+
+        self.model = YOLO(self.model_path)
+        return self.model
+
+    @staticmethod
+    def parking_regions_extraction(json_file):
+        """
+        Extract parking regions from json file.
+
+        Args:
+            json_file (str): file that have all parking slot points
+        """
+
+        with open(json_file, "r") as json_file:
+            return json.load(json_file)
+
+    def process_data(self, json_data, im0, boxes, clss):
+        """
+        Process the model data for parking lot management.
+
+        Args:
+            json_data (str): json data for parking lot management
+            im0 (ndarray): inference image
+            boxes (list): bounding boxes data
+            clss (list): bounding boxes classes list
+        Returns:
+            filled_slots (int): total slots that are filled in parking lot
+            empty_slots (int): total slots that are available in parking lot
+        """
+        annotator = Annotator(im0)
+        total_slots, filled_slots = len(json_data), 0
+        empty_slots = total_slots
+
+        for region in json_data:
+            points = region["points"]
+            points_array = np.array(points, dtype=np.int32).reshape((-1, 1, 2))
+            region_occupied = False
+
+            for box, cls in zip(boxes, clss):
+                x_center = int((box[0] + box[2]) / 2)
+                y_center = int((box[1] + box[3]) / 2)
+                text = f"{self.model.names[int(cls)]}"
+
+                annotator.display_objects_labels(
+                    im0, text, self.txt_color, self.bg_color, x_center, y_center, self.margin
+                )
+                dist = cv2.pointPolygonTest(points_array, (x_center, y_center), False)
+                if dist >= 0:
+                    region_occupied = True
+                    break
+
+            color = self.occupied_region_color if region_occupied else self.available_region_color
+            cv2.polylines(im0, [points_array], isClosed=True, color=color, thickness=2)
+            if region_occupied:
+                filled_slots += 1
+                empty_slots -= 1
+
+        self.labels_dict["Occupancy"] = filled_slots
+        self.labels_dict["Available"] = empty_slots
+
+        annotator.display_analytics(im0, self.labels_dict, self.txt_color, self.bg_color, self.margin)
+
+    def display_frames(self, im0):
+        """
+        Display frame.
+
+        Args:
+            im0 (ndarray): inference image
+        """
+        if self.env_check:
+            cv2.namedWindow(self.window_name)
+            cv2.imshow(self.window_name, im0)
+            # Break Window
+            if cv2.waitKey(1) & 0xFF == ord("q"):
+                return

ultralytics/solutions/queue_management.py

@@ -0,0 +1,187 @@
+# 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 Point, Polygon
+
+
+class QueueManager:
+    """A class to manage the queue management in real-time video stream based on their 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(
+        self,
+        classes_names,
+        reg_pts,
+        line_thickness=2,
+        track_thickness=2,
+        view_img=False,
+        region_color=(255, 0, 255),
+        view_queue_counts=True,
+        draw_tracks=False,
+        count_txt_color=(255, 255, 255),
+        track_color=None,
+        region_thickness=5,
+        fontsize=0.7,
+    ):
+        """
+        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_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
+        """
+        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
+
+        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
+        self.track_color = track_color
+        self.count_txt_color = count_txt_color
+        self.region_thickness = region_thickness
+        self.fontsize = fontsize
+
+    def extract_and_process_tracks(self, tracks):
+        """Extracts and processes tracks for queue management in a video stream."""
+
+        # Annotator Init and queue 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"{self.names[cls]}#{track_id}", 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 if self.track_color else 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
+
+                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)
+
+        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.display_frames()
+
+    def display_frames(self):
+        """Display frame."""
+        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
+            if cv2.waitKey(1) & 0xFF == ord("q"):
+                return
+
+    def process_queue(self, im0, tracks):
+        """
+        Main function to start the queue management 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__":
+    QueueManager()

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()