ultralytics 8.0.238__py3-none-any.whl → 8.0.239__py3-none-any.whl

ultralytics/solutions/ai_gym.py

@@ -26,7 +26,7 @@ class AIGym:
         self.angle = None
         self.count = None
         self.stage = None
-        self.pose_type = 'pushup'
+        self.pose_type = "pushup"
         self.kpts_to_check = None
 
         # Visual Information
@@ -36,13 +36,15 @@ class AIGym:
         # 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'):
+    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:
@@ -72,65 +74,75 @@ class AIGym:
         if frame_count == 1:
             self.count = [0] * len(results[0])
             self.angle = [0] * len(results[0])
-            self.stage = ['-' for _ in 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 == 'pushup' or self.pose_type == '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())
+            if self.pose_type == "pushup" or self.pose_type == "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())
+            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.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':
+                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.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':
+                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.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.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'):
+            cv2.imshow("Ultralytics YOLOv8 AI GYM", self.im0)
+            if cv2.waitKey(1) & 0xFF == ord("q"):
                 return
 
         return self.im0
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     AIGym()

ultralytics/solutions/distance_calculation.py

@@ -41,13 +41,15 @@ class DistanceCalculation:
         # 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)):
+    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.
 
@@ -129,8 +131,9 @@ class DistanceCalculation:
             distance (float): Distance between two centroids
         """
         cv2.rectangle(self.im0, (15, 25), (280, 70), (255, 255, 255), -1)
-        cv2.putText(self.im0, f'Distance : {distance:.2f}m', (20, 55), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2,
-                    cv2.LINE_AA)
+        cv2.putText(
+            self.im0, f"Distance : {distance:.2f}m", (20, 55), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2, cv2.LINE_AA
+        )
         cv2.line(self.im0, self.centroids[0], self.centroids[1], self.line_color, 3)
         cv2.circle(self.im0, self.centroids[0], 6, self.centroid_color, -1)
         cv2.circle(self.im0, self.centroids[1], 6, self.centroid_color, -1)
@@ -179,13 +182,13 @@ class DistanceCalculation:
 
     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)
+        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'):
+        if cv2.waitKey(1) & 0xFF == ord("q"):
             return
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     DistanceCalculation()

ultralytics/solutions/heatmap.py

@@ -8,7 +8,7 @@ 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')
+check_requirements("shapely>=2.0.0")
 
 from shapely.geometry import LineString, Point, Polygon
 
@@ -22,7 +22,7 @@ class Heatmap:
         # Visual information
         self.annotator = None
         self.view_img = False
-        self.shape = 'circle'
+        self.shape = "circle"
 
         # Image information
         self.imw = None
@@ -63,23 +63,25 @@ class Heatmap:
         # 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'):
+    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.
 
@@ -111,20 +113,19 @@ class Heatmap:
 
         # Region and line selection
         if count_reg_pts is not None:
-
             if len(count_reg_pts) == 2:
-                print('Line Counter Initiated.')
+                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.')
+                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')
+                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
@@ -140,10 +141,10 @@ class Heatmap:
         self.shape = shape
 
         # shape of heatmap, if not selected
-        if self.shape not in ['circle', 'rect']:
+        if self.shape not in ["circle", "rect"]:
             print("Unknown shape value provided, 'circle' & 'rect' supported")
-            print('Using Circular shape now')
-            self.shape = 'circle'
+            print("Using Circular shape now")
+            self.shape = "circle"
 
     def extract_results(self, tracks):
         """
@@ -177,27 +178,26 @@ class Heatmap:
         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)
+                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':
+                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
+                    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])])
+                    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
+                    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]
@@ -226,26 +226,26 @@ class Heatmap:
                                 self.in_counts += 1
         else:
             for box, cls in zip(self.boxes, self.clss):
-
-                if self.shape == 'circle':
+                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
+                    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])])
+                    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
+                    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 = 'In Count : ' + f'{self.in_counts}'
-        outcount_label = 'OutCount : ' + f'{self.out_counts}'
+        incount_label = "In Count : " + f"{self.in_counts}"
+        outcount_label = "OutCount : " + f"{self.out_counts}"
 
         # Display counts based on user choice
         counts_label = None
@@ -256,13 +256,15 @@ class Heatmap:
         elif not self.view_out_counts:
             counts_label = incount_label
         else:
-            counts_label = incount_label + ' ' + outcount_label
+            counts_label = 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.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)
 
@@ -273,11 +275,11 @@ class Heatmap:
 
     def display_frames(self):
         """Display frame."""
-        cv2.imshow('Ultralytics Heatmap', self.im0)
+        cv2.imshow("Ultralytics Heatmap", self.im0)
 
-        if cv2.waitKey(1) & 0xFF == ord('q'):
+        if cv2.waitKey(1) & 0xFF == ord("q"):
             return
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     Heatmap()

ultralytics/solutions/object_counter.py

@@ -7,7 +7,7 @@ import cv2
 from ultralytics.utils.checks import check_imshow, check_requirements
 from ultralytics.utils.plotting import Annotator, colors
 
-check_requirements('shapely>=2.0.0')
+check_requirements("shapely>=2.0.0")
 
 from shapely.geometry import LineString, Point, Polygon
 
@@ -56,22 +56,24 @@ class ObjectCounter:
         # 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):
+    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.
 
@@ -101,16 +103,16 @@ class ObjectCounter:
 
         # Region and line selection
         if len(reg_pts) == 2:
-            print('Line Counter Initiated.')
+            print("Line Counter Initiated.")
             self.reg_pts = reg_pts
             self.counting_region = LineString(self.reg_pts)
         elif len(reg_pts) == 4:
-            print('Region Counter Initiated.')
+            print("Region Counter Initiated.")
             self.reg_pts = reg_pts
             self.counting_region = Polygon(self.reg_pts)
         else:
-            print('Invalid Region points provided, region_points can be 2 or 4')
-            print('Using Line Counter Now')
+            print("Invalid Region points provided, region_points can be 2 or 4")
+            print("Using Line Counter Now")
             self.counting_region = LineString(self.reg_pts)
 
         self.names = classes_names
@@ -153,6 +155,7 @@ class ObjectCounter:
             self.selected_point = None
 
     def extract_and_process_tracks(self, tracks):
+        """Extracts and processes tracks for object counting in a video stream."""
         boxes = tracks[0].boxes.xyxy.cpu()
         clss = tracks[0].boxes.cls.cpu().tolist()
         track_ids = tracks[0].boxes.id.int().cpu().tolist()
@@ -163,8 +166,9 @@ class ObjectCounter:
 
         # Extract tracks
         for box, track_id, cls in zip(boxes, track_ids, clss):
-            self.annotator.box_label(box, label=str(track_id) + ':' + self.names[cls],
-                                     color=colors(int(cls), True))  # Draw bounding box
+            self.annotator.box_label(
+                box, label=str(track_id) + ":" + self.names[cls], color=colors(int(cls), True)
+            )  # Draw bounding box
 
             # Draw Tracks
             track_line = self.track_history[track_id]
@@ -174,9 +178,9 @@ class ObjectCounter:
 
             # Draw track trails
             if self.draw_tracks:
-                self.annotator.draw_centroid_and_tracks(track_line,
-                                                        color=self.track_color,
-                                                        track_thickness=self.track_thickness)
+                self.annotator.draw_centroid_and_tracks(
+                    track_line, color=self.track_color, track_thickness=self.track_thickness
+                )
 
             # Count objects
             if len(self.reg_pts) == 4:
@@ -198,8 +202,8 @@ class ObjectCounter:
                         else:
                             self.in_counts += 1
 
-        incount_label = 'In Count : ' + f'{self.in_counts}'
-        outcount_label = 'OutCount : ' + f'{self.out_counts}'
+        incount_label = "In Count : " + f"{self.in_counts}"
+        outcount_label = "OutCount : " + f"{self.out_counts}"
 
         # Display counts based on user choice
         counts_label = None
@@ -210,24 +214,27 @@ class ObjectCounter:
         elif not self.view_out_counts:
             counts_label = incount_label
         else:
-            counts_label = incount_label + ' ' + outcount_label
+            counts_label = 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)
+            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:
-            cv2.namedWindow('Ultralytics YOLOv8 Object Counter')
+            cv2.namedWindow("Ultralytics YOLOv8 Object Counter")
             if len(self.reg_pts) == 4:  # only add mouse event If user drawn region
-                cv2.setMouseCallback('Ultralytics YOLOv8 Object Counter', self.mouse_event_for_region,
-                                     {'region_points': self.reg_pts})
-            cv2.imshow('Ultralytics YOLOv8 Object Counter', self.im0)
+                cv2.setMouseCallback(
+                    "Ultralytics YOLOv8 Object Counter", self.mouse_event_for_region, {"region_points": self.reg_pts}
+                )
+            cv2.imshow("Ultralytics YOLOv8 Object Counter", self.im0)
             # Break Window
-            if cv2.waitKey(1) & 0xFF == ord('q'):
+            if cv2.waitKey(1) & 0xFF == ord("q"):
                 return
 
     def start_counting(self, im0, tracks):
@@ -253,5 +260,5 @@ class ObjectCounter:
         return self.im0
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     ObjectCounter()

ultralytics/solutions/speed_estimation.py

@@ -66,7 +66,7 @@ class SpeedEstimator:
             spdl_dist_thresh (int): Euclidean distance threshold for speed line
         """
         if reg_pts is None:
-            print('Region points not provided, using default values')
+            print("Region points not provided, using default values")
         else:
             self.reg_pts = reg_pts
         self.names = names
@@ -114,8 +114,9 @@ class SpeedEstimator:
             cls (str): object class name
             track (list): tracking history for tracks path drawing
         """
-        speed_label = str(int(
-            self.dist_data[track_id])) + 'km/ph' if track_id in self.dist_data else self.names[int(cls)]
+        speed_label = (
+            str(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)
@@ -132,19 +133,16 @@ class SpeedEstimator:
         """
 
         if self.reg_pts[0][0] < track[-1][0] < self.reg_pts[1][0]:
+            if self.reg_pts[1][1] - self.spdl_dist_thresh < track[-1][1] < self.reg_pts[1][1] + self.spdl_dist_thresh:
+                direction = "known"
 
-            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'
+            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':
+                direction = "unknown"
 
+            if self.trk_previous_times[trk_id] != 0 and direction != "unknown":
                 if trk_id not in self.trk_idslist:
                     self.trk_idslist.append(trk_id)
 
@@ -178,7 +176,6 @@ class SpeedEstimator:
         self.annotator.draw_region(reg_pts=self.reg_pts, color=(255, 0, 0), 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:
@@ -194,10 +191,10 @@ class SpeedEstimator:
 
     def display_frames(self):
         """Display frame."""
-        cv2.imshow('Ultralytics Speed Estimation', self.im0)
-        if cv2.waitKey(1) & 0xFF == ord('q'):
+        cv2.imshow("Ultralytics Speed Estimation", self.im0)
+        if cv2.waitKey(1) & 0xFF == ord("q"):
             return
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     SpeedEstimator()

ultralytics/trackers/__init__.py

@@ -4,4 +4,4 @@ from .bot_sort import BOTSORT
 from .byte_tracker import BYTETracker
 from .track import register_tracker
 
-__all__ = 'register_tracker', 'BOTSORT', 'BYTETracker'  # allow simpler import
+__all__ = "register_tracker", "BOTSORT", "BYTETracker"  # allow simpler import

ultralytics/trackers/basetrack.py

@@ -55,6 +55,7 @@ class BaseTrack:
     _count = 0
 
     def __init__(self):
+        """Initializes a new track with unique ID and foundational tracking attributes."""
         self.track_id = 0
         self.is_activated = False
         self.state = TrackState.New