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

ultralytics/solutions/queue_management.py

@@ -13,49 +13,12 @@ from shapely.geometry import Point, Polygon
 
 
 class QueueManager:
-    """A class to manage the queue management in real-time video stream based on their tracks."""
+    """A class to manage the queue in a real-time video stream based on object 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(
+    def __init__(
         self,
         classes_names,
-        reg_pts,
+        reg_pts=None,
         line_thickness=2,
         track_thickness=2,
         view_img=False,
@@ -68,48 +31,65 @@ class QueueManager:
         fontsize=0.7,
     ):
         """
-        Configures the Counter's image, bounding box line thickness, and counting region points.
+        Initializes the QueueManager with specified parameters for tracking and counting objects.
 
         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
+            classes_names (dict): A dictionary mapping class IDs to class names.
+            reg_pts (list of tuples, optional): Points defining the counting region polygon. Defaults to a predefined
+                rectangle.
+            line_thickness (int, optional): Thickness of the annotation lines. Defaults to 2.
+            track_thickness (int, optional): Thickness of the track lines. Defaults to 2.
+            view_img (bool, optional): Whether to display the image frames. Defaults to False.
+            region_color (tuple, optional): Color of the counting region lines (BGR). Defaults to (255, 0, 255).
+            view_queue_counts (bool, optional): Whether to display the queue counts. Defaults to True.
+            draw_tracks (bool, optional): Whether to draw tracks of the objects. Defaults to False.
+            count_txt_color (tuple, optional): Color of the count text (BGR). Defaults to (255, 255, 255).
+            track_color (tuple, optional): Color of the tracks. If None, different colors will be used for different
+                tracks. Defaults to None.
+            region_thickness (int, optional): Thickness of the counting region lines. Defaults to 5.
+            fontsize (float, optional): Font size for the text annotations. Defaults to 0.7.
         """
+
+        # Mouse events state
+        self.is_drawing = False
+        self.selected_point = None
+
+        # Region & Line Information
+        self.reg_pts = reg_pts if reg_pts is not None else [(20, 60), (20, 680), (1120, 680), (1120, 60)]
+        self.counting_region = (
+            Polygon(self.reg_pts) if len(self.reg_pts) >= 3 else Polygon([(20, 60), (20, 680), (1120, 680), (1120, 60)])
+        )
+        self.region_color = region_color
+        self.region_thickness = region_thickness
+
+        # Image and annotation Information
+        self.im0 = None
         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
+        self.fontsize = fontsize
 
-        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  # Class names
+        self.annotator = None  # Annotator
+        self.window_name = "Ultralytics YOLOv8 Queue Manager"
 
-        self.names = classes_names
-        self.track_color = track_color
+        # Object counting Information
+        self.counts = 0
         self.count_txt_color = count_txt_color
-        self.region_thickness = region_thickness
-        self.fontsize = fontsize
+
+        # Tracks info
+        self.track_history = defaultdict(list)
+        self.track_thickness = track_thickness
+        self.draw_tracks = draw_tracks
+        self.track_color = track_color
+
+        # Check if environment supports imshow
+        self.env_check = check_imshow(warn=True)
 
     def extract_and_process_tracks(self, tracks):
         """Extracts and processes tracks for queue management in a video stream."""
 
-        # Annotator Init and queue region drawing
+        # Initialize annotator and draw the queue region
         self.annotator = Annotator(self.im0, self.tf, self.names)
 
         if tracks[0].boxes.id is not None:
@@ -122,48 +102,48 @@ class QueueManager:
                 # Draw bounding box
                 self.annotator.box_label(box, label=f"{self.names[cls]}#{track_id}", color=colors(int(track_id), True))
 
-                # Draw Tracks
+                # Update track history
                 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
+                # Draw track trails if enabled
                 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),
+                        color=self.track_color or 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
 
+                # Check if the object is inside the counting region
                 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)
-
+        # Display queue counts
+        label = f"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.counts = 0  # Reset counts after displaying
         self.display_frames()
 
     def display_frames(self):
-        """Display frame."""
+        """Displays the current frame with annotations."""
         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
+            # Close window on 'q' key press
             if cv2.waitKey(1) & 0xFF == ord("q"):
                 return
 
@@ -175,13 +155,14 @@ class QueueManager:
             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
+        self.im0 = im0  # Store the current frame
+        self.extract_and_process_tracks(tracks)  # Extract and process tracks
 
         if self.view_img:
-            self.display_frames()
+            self.display_frames()  # Display the frame if enabled
         return self.im0
 
 
 if __name__ == "__main__":
-    QueueManager()
+    classes_names = {0: "person", 1: "car"}  # example class names
+    queue_manager = QueueManager(classes_names)

ultralytics/solutions/speed_estimation.py

@@ -11,73 +11,52 @@ 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."""
+    """A class to estimate the speed of objects in a 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."""
+    def __init__(self, names, reg_pts=None, view_img=False, line_thickness=2, region_thickness=5, spdl_dist_thresh=10):
+        """
+        Initializes the SpeedEstimator with the given parameters.
 
-        # Visual & im0 information
+        Args:
+            names (dict): Dictionary of class names.
+            reg_pts (list, optional): List of region points for speed estimation. Defaults to [(20, 400), (1260, 400)].
+            view_img (bool, optional): Whether to display the image with annotations. Defaults to False.
+            line_thickness (int, optional): Thickness of the lines for drawing boxes and tracks. Defaults to 2.
+            region_thickness (int, optional): Thickness of the region lines. Defaults to 5.
+            spdl_dist_thresh (int, optional): Distance threshold for speed calculation. Defaults to 10.
+        """
+        # Visual & image information
         self.im0 = None
         self.annotator = None
-        self.view_img = False
+        self.view_img = view_img
 
         # Region information
-        self.reg_pts = [(20, 400), (1260, 400)]
-        self.region_thickness = 3
+        self.reg_pts = reg_pts if reg_pts is not None else [(20, 400), (1260, 400)]
+        self.region_thickness = region_thickness
 
-        # Predict/track information
+        # Tracking information
         self.clss = None
-        self.names = None
+        self.names = names
         self.boxes = None
         self.trk_ids = None
         self.trk_pts = None
-        self.line_thickness = 2
+        self.line_thickness = line_thickness
         self.trk_history = defaultdict(list)
 
-        # Speed estimator information
+        # Speed estimation information
         self.current_time = 0
         self.dist_data = {}
         self.trk_idslist = []
-        self.spdl_dist_thresh = 10
+        self.spdl_dist_thresh = spdl_dist_thresh
         self.trk_previous_times = {}
         self.trk_previous_points = {}
 
-        # Check if environment support imshow
+        # Check if the environment supports 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.
+        Extracts results from the provided tracking data.
 
         Args:
             tracks (list): List of tracks obtained from the object tracking process.
@@ -88,11 +67,14 @@ class SpeedEstimator:
 
     def store_track_info(self, track_id, box):
         """
-        Store track data.
+        Stores track data.
 
         Args:
-            track_id (int): object track id.
-            box (list): object bounding box data
+            track_id (int): Object track id.
+            box (list): Object bounding box data.
+
+        Returns:
+            (list): Updated tracking history for the given track_id.
         """
         track = self.trk_history[track_id]
         bbox_center = (float((box[0] + box[2]) / 2), float((box[1] + box[3]) / 2))
@@ -106,43 +88,39 @@ class SpeedEstimator:
 
     def plot_box_and_track(self, track_id, box, cls, track):
         """
-        Plot track and bounding box.
+        Plots 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
+            track_id (int): Object track id.
+            box (list): Object bounding box data.
+            cls (str): Object class name.
+            track (list): Tracking history for drawing tracks path.
         """
-        speed_label = f"{int(self.dist_data[track_id])}km/ph" if track_id in self.dist_data else self.names[int(cls)]
+        speed_label = f"{int(self.dist_data[track_id])} km/h" 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.
+        Calculates the speed of an object.
 
         Args:
-            trk_id (int): object track id.
-            track (list): tracking history for tracks path drawing
+            trk_id (int): Object track id.
+            track (list): Tracking history for drawing tracks path.
         """
-
         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:
+        if self.trk_previous_times.get(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]
@@ -156,21 +134,24 @@ class SpeedEstimator:
 
     def estimate_speed(self, im0, tracks, region_color=(255, 0, 0)):
         """
-        Calculate object based on tracking data.
+        Estimates the speed of objects based on tracking data.
 
         Args:
-            im0 (nd array): Image
+            im0 (ndarray): Image.
             tracks (list): List of tracks obtained from the object tracking process.
-            region_color (tuple): Color to use when drawing regions.
+            region_color (tuple, optional): Color to use when drawing regions. Defaults to (255, 0, 0).
+
+        Returns:
+            (ndarray): The image with annotated boxes and tracks.
         """
         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.extract_tracks(tracks)
+        self.annotator = Annotator(self.im0, line_width=self.line_thickness)
         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):
@@ -188,11 +169,12 @@ class SpeedEstimator:
         return im0
 
     def display_frames(self):
-        """Display frame."""
+        """Displays the current frame."""
         cv2.imshow("Ultralytics Speed Estimation", self.im0)
         if cv2.waitKey(1) & 0xFF == ord("q"):
             return
 
 
 if __name__ == "__main__":
-    SpeedEstimator()
+    names = {0: "person", 1: "car"}  # example class names
+    speed_estimator = SpeedEstimator(names)

ultralytics/utils/__init__.py

@@ -515,7 +515,7 @@ def is_online() -> bool:
         import socket
 
         for dns in ("1.1.1.1", "8.8.8.8"):  # check Cloudflare and Google DNS
-            socket.create_connection(address=(dns, 80), timeout=1.0).close()
+            socket.create_connection(address=(dns, 80), timeout=2.0).close()
             return True
     return False
 

ultralytics/utils/callbacks/wb.py

@@ -100,7 +100,7 @@ def _plot_curve(
 
 def _log_plots(plots, step):
     """Logs plots from the input dictionary if they haven't been logged already at the specified step."""
-    for name, params in plots.items():
+    for name, params in plots.copy().items():  # shallow copy to prevent plots dict changing during iteration
         timestamp = params["timestamp"]
         if _processed_plots.get(name) != timestamp:
             wb.run.log({name.stem: wb.Image(str(name))}, step=step)

ultralytics/utils/checks.py

@@ -391,7 +391,7 @@ def check_requirements(requirements=ROOT.parent / "requirements.txt", exclude=()
                 t = time.time()
                 assert ONLINE, "AutoUpdate skipped (offline)"
                 with Retry(times=2, delay=1):  # run up to 2 times with 1-second retry delay
-                    LOGGER.info(subprocess.check_output(f"pip install --no-cache {s} {cmds}", shell=True).decode())
+                    LOGGER.info(subprocess.check_output(f"pip install --no-cache-dir {s} {cmds}", shell=True).decode())
                 dt = time.time() - t
                 LOGGER.info(
                     f"{prefix} AutoUpdate success ✅ {dt:.1f}s, installed {n} package{'s' * (n > 1)}: {pkgs}\n"
@@ -535,7 +535,8 @@ def check_imshow(warn=False):
     """Check if environment supports image displays."""
     try:
         if LINUX:
-            assert "DISPLAY" in os.environ and not IS_DOCKER and not IS_COLAB and not IS_KAGGLE
+            assert not IS_COLAB and not IS_KAGGLE
+            assert "DISPLAY" in os.environ, "The DISPLAY environment variable isn't set."
         cv2.imshow("test", np.zeros((8, 8, 3), dtype=np.uint8))  # show a small 8-pixel image
         cv2.waitKey(1)
         cv2.destroyAllWindows()