ultralytics 8.3.4__py3-none-any.whl → 8.3.6__py3-none-any.whl

ultralytics/solutions/object_counter.py

@@ -1,243 +1,131 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
-from collections import defaultdict
+from shapely.geometry import LineString, Point
 
-import cv2
-
-from ultralytics.utils.checks import check_imshow, check_requirements
+from ultralytics.solutions.solutions import BaseSolution  # Import a parent class
 from ultralytics.utils.plotting import Annotator, colors
 
-check_requirements("shapely>=2.0.0")
 
-from shapely.geometry import LineString, Point, Polygon
+class ObjectCounter(BaseSolution):
+    """A class to manage the counting of objects in a real-time video stream based on their tracks."""
 
+    def __init__(self, **kwargs):
+        """Initialization function for Count class, a child class of BaseSolution class, can be used for counting the
+        objects.
+        """
+        super().__init__(**kwargs)
 
-class ObjectCounter:
-    """A class to manage the counting of objects in a real-time video stream based on their tracks."""
+        self.in_count = 0  # Counter for objects moving inward
+        self.out_count = 0  # Counter for objects moving outward
+        self.counted_ids = []  # List of IDs of objects that have been counted
+        self.classwise_counts = {}  # Dictionary for counts, categorized by object class
+        self.region_initialized = False  # Bool variable for region initialization
 
-    def __init__(
-        self,
-        names,
-        reg_pts=None,
-        line_thickness=2,
-        view_img=False,
-        view_in_counts=True,
-        view_out_counts=True,
-        draw_tracks=False,
-    ):
+        self.show_in = self.CFG["show_in"]
+        self.show_out = self.CFG["show_out"]
+
+    def count_objects(self, track_line, box, track_id, prev_position, cls):
         """
-        Initializes the ObjectCounter with various tracking and counting parameters.
+        Helper function to count objects within a polygonal region.
 
         Args:
-            names (dict): Dictionary of class names.
-            reg_pts (list): List of points defining the counting region.
-            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 in counts on the video stream.
-            view_out_counts (bool): Flag to control whether to display the out counts on the video stream.
-            draw_tracks (bool): Flag to control whether to draw the object tracks.
+            track_line (dict): last 30 frame track record
+            box (list): Bounding box data for specific track in current frame
+            track_id (int): track ID of the object
+            prev_position (tuple): last frame position coordinates of the track
+            cls (int): Class index for classwise count updates
         """
-        # Mouse events
-        self.is_drawing = False
-        self.selected_point = None
-
-        # Region & Line Information
-        self.reg_pts = [(20, 400), (1260, 400)] if reg_pts is None else reg_pts
-        self.counting_region = None
-
-        # Image and annotation Information
-        self.im0 = None
-        self.tf = line_thickness
-        self.view_img = view_img
-        self.view_in_counts = view_in_counts
-        self.view_out_counts = view_out_counts
-
-        self.names = names  # Classes names
-        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 = {}
-
-        # Tracks info
-        self.track_history = defaultdict(list)
-        self.draw_tracks = draw_tracks
-
-        # Check if environment supports imshow
-        self.env_check = check_imshow(warn=True)
-
-        # Initialize counting region
-        if len(self.reg_pts) == 2:
-            print("Line Counter Initiated.")
-            self.counting_region = LineString(self.reg_pts)
-        elif len(self.reg_pts) >= 3:
-            print("Polygon Counter Initiated.")
-            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)
-
-        # Define the counting line segment
-        self.counting_line_segment = LineString(
-            [
-                (self.reg_pts[0][0], self.reg_pts[0][1]),
-                (self.reg_pts[1][0], self.reg_pts[1][1]),
-            ]
-        )
-
-    def mouse_event_for_region(self, event, x, y, flags, params):
+        if prev_position is None or track_id in self.counted_ids:
+            return
+
+        centroid = self.r_s.centroid
+        dx = (box[0] - prev_position[0]) * (centroid.x - prev_position[0])
+        dy = (box[1] - prev_position[1]) * (centroid.y - prev_position[1])
+
+        if len(self.region) >= 3 and self.r_s.contains(Point(track_line[-1])):
+            self.counted_ids.append(track_id)
+            # For polygon region
+            if dx > 0:
+                self.in_count += 1
+                self.classwise_counts[self.names[cls]]["IN"] += 1
+            else:
+                self.out_count += 1
+                self.classwise_counts[self.names[cls]]["OUT"] += 1
+
+        elif len(self.region) < 3 and LineString([prev_position, box[:2]]).intersects(self.l_s):
+            self.counted_ids.append(track_id)
+            # For linear region
+            if dx > 0 and dy > 0:
+                self.in_count += 1
+                self.classwise_counts[self.names[cls]]["IN"] += 1
+            else:
+                self.out_count += 1
+                self.classwise_counts[self.names[cls]]["OUT"] += 1
+
+    def store_classwise_counts(self, cls):
         """
-        Handles mouse events for defining and moving the counting region in a real-time video stream.
+        Initialize class-wise counts if not already present.
 
         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 associated event flags (e.g., cv2.EVENT_FLAG_CTRLKEY,  cv2.EVENT_FLAG_SHIFTKEY, etc.).
-            params (dict): Additional parameters for the function.
+            cls (int): Class index for classwise count updates
+        """
+        if self.names[cls] not in self.classwise_counts:
+            self.classwise_counts[self.names[cls]] = {"IN": 0, "OUT": 0}
+
+    def display_counts(self, im0):
         """
-        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
-        annotator = Annotator(self.im0, self.tf, self.names)
-
-        # Draw region or line
-        annotator.draw_region(reg_pts=self.reg_pts, color=(104, 0, 123), thickness=self.tf * 2)
-
-        # Extract tracks for OBB or object detection
-        track_data = tracks[0].obb or tracks[0].boxes
-
-        if track_data and track_data.id is not None:
-            boxes = track_data.xyxy.cpu()
-            clss = track_data.cls.cpu().tolist()
-            track_ids = track_data.id.int().cpu().tolist()
-
-            # Extract tracks
-            for box, track_id, cls in zip(boxes, track_ids, clss):
-                # Draw bounding box
-                annotator.box_label(box, label=self.names[cls], 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:
-                    annotator.draw_centroid_and_tracks(
-                        track_line,
-                        color=colors(int(track_id), True),
-                        track_thickness=self.tf,
-                    )
-
-                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
-                        and LineString([(prev_position[0], prev_position[1]), (box[0], box[1])]).intersects(
-                            self.counting_line_segment
-                        )
-                    ):
-                        self.count_ids.append(track_id)
-
-                        # Determine the direction of movement (IN or OUT)
-                        dx = (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0])
-                        dy = (box[1] - prev_position[1]) * (self.counting_region.centroid.y - prev_position[1])
-                        if dx > 0 and dy > 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']}"
+        Helper function to display object counts on the frame.
+
+        Args:
+            im0 (ndarray): The input image or frame
+        """
+        labels_dict = {
+            str.capitalize(key): f"{'IN ' + str(value['IN']) if self.show_in else ''} "
+            f"{'OUT ' + str(value['OUT']) if self.show_out else ''}".strip()
+            for key, value in self.classwise_counts.items()
+            if value["IN"] != 0 or value["OUT"] != 0
+        }
 
         if labels_dict:
-            annotator.display_analytics(self.im0, labels_dict, (104, 31, 17), (255, 255, 255), 10)
-
-    def display_frames(self):
-        """Displays the current frame with annotations and regions in a window."""
-        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):
+            self.annotator.display_analytics(im0, labels_dict, (104, 31, 17), (255, 255, 255), 10)
+
+    def count(self, im0):
         """
-        Main function to start the object counting process.
+        Processes input data (frames or object tracks) and updates counts.
 
         Args:
-            im0 (ndarray): Current frame from the video stream.
-            tracks (list): List of tracks obtained from the object tracking process.
+            im0 (ndarray): The input image that will be used for processing
+        Returns
+            im0 (ndarray): The processed image for more usage
         """
-        self.im0 = im0  # store image
-        self.extract_and_process_tracks(tracks)  # draw region even if no objects
+        if not self.region_initialized:
+            self.initialize_region()
+            self.region_initialized = True
+
+        self.annotator = Annotator(im0, line_width=self.line_width)  # Initialize annotator
+        self.extract_tracks(im0)  # Extract tracks
+
+        self.annotator.draw_region(
+            reg_pts=self.region, color=(104, 0, 123), thickness=self.line_width * 2
+        )  # Draw region
+
+        # Iterate over bounding boxes, track ids and classes index
+        for box, track_id, cls in zip(self.boxes, self.track_ids, self.clss):
+            # Draw bounding box and counting region
+            self.annotator.box_label(box, label=self.names[cls], color=colors(track_id, True))
+            self.store_tracking_history(track_id, box)  # Store track history
+            self.store_classwise_counts(cls)  # store classwise counts in dict
+
+            # Draw centroid of objects
+            self.annotator.draw_centroid_and_tracks(
+                self.track_line, color=colors(int(track_id), True), track_thickness=self.line_width
+            )
 
-        if self.view_img:
-            self.display_frames()
-        return self.im0
+            # store previous position of track for object counting
+            prev_position = self.track_history[track_id][-2] if len(self.track_history[track_id]) > 1 else None
+            self.count_objects(self.track_line, box, track_id, prev_position, cls)  # Perform object counting
 
+        self.display_counts(im0)  # Display the counts on the frame
+        self.display_output(im0)  # display output with base class function
 
-if __name__ == "__main__":
-    classes_names = {0: "person", 1: "car"}  # example class names
-    ObjectCounter(classes_names)
+        return im0  # return output image for more usage

ultralytics/solutions/solutions.py

@@ -0,0 +1,93 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+
+from collections import defaultdict
+from pathlib import Path
+
+import cv2
+
+from ultralytics import YOLO
+from ultralytics.utils import LOGGER, yaml_load
+from ultralytics.utils.checks import check_imshow, check_requirements
+
+check_requirements("shapely>=2.0.0")
+from shapely.geometry import LineString, Polygon
+
+DEFAULT_SOL_CFG_PATH = Path(__file__).resolve().parents[1] / "cfg/solutions/default.yaml"
+
+
+class BaseSolution:
+    """A class to manage all the Ultralytics Solutions: https://docs.ultralytics.com/solutions/."""
+
+    def __init__(self, **kwargs):
+        """
+        Base initializer for all solutions.
+
+        Child classes should call this with necessary parameters.
+        """
+        # Load config and update with args
+        self.CFG = yaml_load(DEFAULT_SOL_CFG_PATH)
+        self.CFG.update(kwargs)
+        LOGGER.info(f"Ultralytics Solutions: ✅ {self.CFG}")
+
+        self.region = self.CFG["region"]  # Store region data for other classes usage
+        self.line_width = self.CFG["line_width"]  # Store line_width for usage
+
+        # Load Model and store classes names
+        self.model = YOLO(self.CFG["model"])
+        self.names = self.model.names
+
+        # Initialize environment and region setup
+        self.env_check = check_imshow(warn=True)
+        self.track_history = defaultdict(list)
+
+    def extract_tracks(self, im0):
+        """
+        Apply object tracking and extract tracks.
+
+        Args:
+            im0 (ndarray): The input image or frame
+        """
+        self.tracks = self.model.track(source=im0, persist=True, classes=self.CFG["classes"])
+
+        # Extract tracks for OBB or object detection
+        self.track_data = self.tracks[0].obb or self.tracks[0].boxes
+
+        if self.track_data and self.track_data.id is not None:
+            self.boxes = self.track_data.xyxy.cpu()
+            self.clss = self.track_data.cls.cpu().tolist()
+            self.track_ids = self.track_data.id.int().cpu().tolist()
+        else:
+            LOGGER.warning("WARNING ⚠️ no tracks found!")
+            self.boxes, self.clss, self.track_ids = [], [], []
+
+    def store_tracking_history(self, track_id, box):
+        """
+        Store object tracking history.
+
+        Args:
+            track_id (int): The track ID of the object
+            box (list): Bounding box coordinates of the object
+        """
+        # Store tracking history
+        self.track_line = self.track_history[track_id]
+        self.track_line.append(((box[0] + box[2]) / 2, (box[1] + box[3]) / 2))
+        if len(self.track_line) > 30:
+            self.track_line.pop(0)
+
+    def initialize_region(self):
+        """Initialize the counting region and line segment based on config."""
+        self.region = [(20, 400), (1260, 400)] if self.region is None else self.region
+        self.r_s = Polygon(self.region) if len(self.region) >= 3 else LineString(self.region)
+        self.l_s = LineString([(self.region[0][0], self.region[0][1]), (self.region[1][0], self.region[1][1])])
+
+    def display_output(self, im0):
+        """
+        Display the results of the processing, which could involve showing frames, printing counts, or saving results.
+
+        Args:
+            im0 (ndarray): The input image or frame
+        """
+        if self.CFG.get("show") and self.env_check:
+            cv2.imshow("Ultralytics Solutions", im0)
+            if cv2.waitKey(1) & 0xFF == ord("q"):
+                return

ultralytics/utils/__init__.py

@@ -61,8 +61,8 @@ HELP_MSG = """
         from ultralytics import YOLO
 
         # Load a model
-        model = YOLO("yolov8n.yaml")  # build a new model from scratch
-        model = YOLO("yolov8n.pt")  # load a pretrained model (recommended for training)
+        model = YOLO("yolo11n.yaml")  # build a new model from scratch
+        model = YOLO("yolo11n.pt")  # load a pretrained model (recommended for training)
 
         # Use the model
         results = model.train(data="coco8.yaml", epochs=3)  # train the model
@@ -77,21 +77,21 @@ HELP_MSG = """
             yolo TASK MODE ARGS
 
             Where   TASK (optional) is one of [detect, segment, classify, pose, obb]
-                    MODE (required) is one of [train, val, predict, export, benchmark]
+                    MODE (required) is one of [train, val, predict, export, track, benchmark]
                     ARGS (optional) are any number of custom "arg=value" pairs like "imgsz=320" that override defaults.
                         See all ARGS at https://docs.ultralytics.com/usage/cfg or with "yolo cfg"
 
         - Train a detection model for 10 epochs with an initial learning_rate of 0.01
-            yolo detect train data=coco8.yaml model=yolov8n.pt epochs=10 lr0=0.01
+            yolo detect train data=coco8.yaml model=yolo11n.pt epochs=10 lr0=0.01
 
         - Predict a YouTube video using a pretrained segmentation model at image size 320:
-            yolo segment predict model=yolov8n-seg.pt source='https://youtu.be/LNwODJXcvt4' imgsz=320
+            yolo segment predict model=yolo11n-seg.pt source='https://youtu.be/LNwODJXcvt4' imgsz=320
 
         - Val a pretrained detection model at batch-size 1 and image size 640:
-            yolo detect val model=yolov8n.pt data=coco8.yaml batch=1 imgsz=640
+            yolo detect val model=yolo11n.pt data=coco8.yaml batch=1 imgsz=640
 
-        - Export a YOLOv8n classification model to ONNX format at image size 224 by 128 (no TASK required)
-            yolo export model=yolov8n-cls.pt format=onnx imgsz=224,128
+        - Export a YOLO11n classification model to ONNX format at image size 224 by 128 (no TASK required)
+            yolo export model=yolo11n-cls.pt format=onnx imgsz=224,128
 
         - Run special commands:
             yolo help
@@ -989,55 +989,56 @@ def set_sentry():
     Additionally, the function sets custom tags and user information for Sentry events.
     """
     if (
-        SETTINGS["sync"]
-        and RANK in {-1, 0}
-        and Path(ARGV[0]).name == "yolo"
-        and not TESTS_RUNNING
-        and ONLINE
-        and IS_PIP_PACKAGE
-        and not IS_GIT_DIR
+        not SETTINGS["sync"]
+        or RANK not in {-1, 0}
+        or Path(ARGV[0]).name != "yolo"
+        or TESTS_RUNNING
+        or not ONLINE
+        or not IS_PIP_PACKAGE
+        or IS_GIT_DIR
     ):
-        # If sentry_sdk package is not installed then return and do not use Sentry
-        try:
-            import sentry_sdk  # noqa
-        except ImportError:
-            return
-
-        def before_send(event, hint):
-            """
-            Modify the event before sending it to Sentry based on specific exception types and messages.
+        return
+    # If sentry_sdk package is not installed then return and do not use Sentry
+    try:
+        import sentry_sdk  # noqa
+    except ImportError:
+        return
+
+    def before_send(event, hint):
+        """
+        Modify the event before sending it to Sentry based on specific exception types and messages.
 
-            Args:
-                event (dict): The event dictionary containing information about the error.
-                hint (dict): A dictionary containing additional information about the error.
+        Args:
+            event (dict): The event dictionary containing information about the error.
+            hint (dict): A dictionary containing additional information about the error.
 
-            Returns:
-                dict: The modified event or None if the event should not be sent to Sentry.
-            """
-            if "exc_info" in hint:
-                exc_type, exc_value, _ = hint["exc_info"]
-                if exc_type in {KeyboardInterrupt, FileNotFoundError} or "out of memory" in str(exc_value):
-                    return None  # do not send event
-
-            event["tags"] = {
-                "sys_argv": ARGV[0],
-                "sys_argv_name": Path(ARGV[0]).name,
-                "install": "git" if IS_GIT_DIR else "pip" if IS_PIP_PACKAGE else "other",
-                "os": ENVIRONMENT,
-            }
-            return event
-
-        sentry_sdk.init(
-            dsn="https://888e5a0778212e1d0314c37d4b9aae5d@o4504521589325824.ingest.us.sentry.io/4504521592406016",
-            debug=False,
-            auto_enabling_integrations=False,
-            traces_sample_rate=1.0,
-            release=__version__,
-            environment="production",  # 'dev' or 'production'
-            before_send=before_send,
-            ignore_errors=[KeyboardInterrupt, FileNotFoundError],
-        )
-        sentry_sdk.set_user({"id": SETTINGS["uuid"]})  # SHA-256 anonymized UUID hash
+        Returns:
+            dict: The modified event or None if the event should not be sent to Sentry.
+        """
+        if "exc_info" in hint:
+            exc_type, exc_value, _ = hint["exc_info"]
+            if exc_type in {KeyboardInterrupt, FileNotFoundError} or "out of memory" in str(exc_value):
+                return None  # do not send event
+
+        event["tags"] = {
+            "sys_argv": ARGV[0],
+            "sys_argv_name": Path(ARGV[0]).name,
+            "install": "git" if IS_GIT_DIR else "pip" if IS_PIP_PACKAGE else "other",
+            "os": ENVIRONMENT,
+        }
+        return event
+
+    sentry_sdk.init(
+        dsn="https://888e5a0778212e1d0314c37d4b9aae5d@o4504521589325824.ingest.us.sentry.io/4504521592406016",
+        debug=False,
+        auto_enabling_integrations=False,
+        traces_sample_rate=1.0,
+        release=__version__,
+        environment="production",  # 'dev' or 'production'
+        before_send=before_send,
+        ignore_errors=[KeyboardInterrupt, FileNotFoundError],
+    )
+    sentry_sdk.set_user({"id": SETTINGS["uuid"]})  # SHA-256 anonymized UUID hash
 
 
 class JSONDict(dict):