dgenerate-ultralytics-headless 8.3.253__py3-none-any.whl

ultralytics/solutions/instance_segmentation.py

@@ -0,0 +1,86 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from typing import Any
+
+from ultralytics.engine.results import Results
+from ultralytics.solutions.solutions import BaseSolution, SolutionResults
+
+
+class InstanceSegmentation(BaseSolution):
+    """A class to manage instance segmentation in images or video streams.
+
+    This class extends the BaseSolution class and provides functionality for performing instance segmentation, including
+    drawing segmented masks with bounding boxes and labels.
+
+    Attributes:
+        model (str): The segmentation model to use for inference.
+        line_width (int): Width of the bounding box and text lines.
+        names (dict[int, str]): Dictionary mapping class indices to class names.
+        clss (list[int]): List of detected class indices.
+        track_ids (list[int]): List of track IDs for detected instances.
+        masks (list[np.ndarray]): List of segmentation masks for detected instances.
+        show_conf (bool): Whether to display confidence scores.
+        show_labels (bool): Whether to display class labels.
+        show_boxes (bool): Whether to display bounding boxes.
+
+    Methods:
+        process: Process the input image to perform instance segmentation and annotate results.
+        extract_tracks: Extract tracks including bounding boxes, classes, and masks from model predictions.
+
+    Examples:
+        >>> segmenter = InstanceSegmentation()
+        >>> frame = cv2.imread("frame.jpg")
+        >>> results = segmenter.process(frame)
+        >>> print(f"Total segmented instances: {results.total_tracks}")
+    """
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Initialize the InstanceSegmentation class for detecting and annotating segmented instances.
+
+        Args:
+            **kwargs (Any): Keyword arguments passed to the BaseSolution parent class including:
+                - model (str): Model name or path, defaults to "yolo11n-seg.pt".
+        """
+        kwargs["model"] = kwargs.get("model", "yolo11n-seg.pt")
+        super().__init__(**kwargs)
+
+        self.show_conf = self.CFG.get("show_conf", True)
+        self.show_labels = self.CFG.get("show_labels", True)
+        self.show_boxes = self.CFG.get("show_boxes", True)
+
+    def process(self, im0) -> SolutionResults:
+        """Perform instance segmentation on the input image and annotate the results.
+
+        Args:
+            im0 (np.ndarray): The input image for segmentation.
+
+        Returns:
+            (SolutionResults): Object containing the annotated image and total number of tracked instances.
+
+        Examples:
+            >>> segmenter = InstanceSegmentation()
+            >>> frame = cv2.imread("image.jpg")
+            >>> summary = segmenter.process(frame)
+            >>> print(summary)
+        """
+        self.extract_tracks(im0)  # Extract tracks (bounding boxes, classes, and masks)
+        self.masks = getattr(self.tracks, "masks", None)
+
+        # Iterate over detected classes, track IDs, and segmentation masks
+        if self.masks is None:
+            self.LOGGER.warning("No masks detected! Ensure you're using a supported Ultralytics segmentation model.")
+            plot_im = im0
+        else:
+            results = Results(im0, path=None, names=self.names, boxes=self.track_data.data, masks=self.masks.data)
+            plot_im = results.plot(
+                line_width=self.line_width,
+                boxes=self.show_boxes,
+                conf=self.show_conf,
+                labels=self.show_labels,
+                color_mode="instance",
+            )
+
+        self.display_output(plot_im)  # Display the annotated output using the base class function
+
+        # Return SolutionResults
+        return SolutionResults(plot_im=plot_im, total_tracks=len(self.track_ids))

ultralytics/solutions/object_blurrer.py

@@ -0,0 +1,89 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from typing import Any
+
+import cv2
+
+from ultralytics.solutions.solutions import BaseSolution, SolutionAnnotator, SolutionResults
+from ultralytics.utils import LOGGER
+from ultralytics.utils.plotting import colors
+
+
+class ObjectBlurrer(BaseSolution):
+    """A class to manage the blurring of detected objects in a real-time video stream.
+
+    This class extends the BaseSolution class and provides functionality for blurring objects based on detected bounding
+    boxes. The blurred areas are updated directly in the input image, allowing for privacy preservation or other effects.
+
+    Attributes:
+        blur_ratio (int): The intensity of the blur effect applied to detected objects (higher values create more blur).
+        iou (float): Intersection over Union threshold for object detection.
+        conf (float): Confidence threshold for object detection.
+
+    Methods:
+        process: Apply a blurring effect to detected objects in the input image.
+        extract_tracks: Extract tracking information from detected objects.
+        display_output: Display the processed output image.
+
+    Examples:
+        >>> blurrer = ObjectBlurrer()
+        >>> frame = cv2.imread("frame.jpg")
+        >>> processed_results = blurrer.process(frame)
+        >>> print(f"Total blurred objects: {processed_results.total_tracks}")
+    """
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Initialize the ObjectBlurrer class for applying a blur effect to objects detected in video streams or images.
+
+        Args:
+            **kwargs (Any): Keyword arguments passed to the parent class and for configuration including:
+                - blur_ratio (float): Intensity of the blur effect (0.1-1.0, default=0.5).
+        """
+        super().__init__(**kwargs)
+        blur_ratio = self.CFG["blur_ratio"]
+        if blur_ratio < 0.1:
+            LOGGER.warning("blur ratio cannot be less than 0.1, updating it to default value 0.5")
+            blur_ratio = 0.5
+        self.blur_ratio = int(blur_ratio * 100)
+
+    def process(self, im0) -> SolutionResults:
+        """Apply a blurring effect to detected objects in the input image.
+
+        This method extracts tracking information, applies blur to regions corresponding to detected objects, and
+        annotates the image with bounding boxes.
+
+        Args:
+            im0 (np.ndarray): The input image containing detected objects.
+
+        Returns:
+            (SolutionResults): Object containing the processed image and number of tracked objects.
+                - plot_im (np.ndarray): The annotated output image with blurred objects.
+                - total_tracks (int): The total number of tracked objects in the frame.
+
+        Examples:
+            >>> blurrer = ObjectBlurrer()
+            >>> frame = cv2.imread("image.jpg")
+            >>> results = blurrer.process(frame)
+            >>> print(f"Blurred {results.total_tracks} objects")
+        """
+        self.extract_tracks(im0)  # Extract tracks
+        annotator = SolutionAnnotator(im0, self.line_width)
+
+        # Iterate over bounding boxes and classes
+        for box, cls, conf in zip(self.boxes, self.clss, self.confs):
+            # Crop and blur the detected object
+            blur_obj = cv2.blur(
+                im0[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])],
+                (self.blur_ratio, self.blur_ratio),
+            )
+            # Update the blurred area in the original image
+            im0[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])] = blur_obj
+            annotator.box_label(
+                box, label=self.adjust_box_label(cls, conf), color=colors(cls, True)
+            )  # Annotate bounding box
+
+        plot_im = annotator.result()
+        self.display_output(plot_im)  # Display the output using the base class function
+
+        # Return a SolutionResults
+        return SolutionResults(plot_im=plot_im, total_tracks=len(self.track_ids))

ultralytics/solutions/object_counter.py

@@ -0,0 +1,190 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import Any
+
+from ultralytics.solutions.solutions import BaseSolution, SolutionAnnotator, SolutionResults
+from ultralytics.utils.plotting import colors
+
+
+class ObjectCounter(BaseSolution):
+    """A class to manage the counting of objects in a real-time video stream based on their tracks.
+
+    This class extends the BaseSolution class and provides functionality for counting objects moving in and out of a
+    specified region in a video stream. It supports both polygonal and linear regions for counting.
+
+    Attributes:
+        in_count (int): Counter for objects moving inward.
+        out_count (int): Counter for objects moving outward.
+        counted_ids (list[int]): List of IDs of objects that have been counted.
+        classwise_count (dict[str, dict[str, int]]): Dictionary for counts, categorized by object class.
+        region_initialized (bool): Flag indicating whether the counting region has been initialized.
+        show_in (bool): Flag to control display of inward count.
+        show_out (bool): Flag to control display of outward count.
+        margin (int): Margin for background rectangle size to display counts properly.
+
+    Methods:
+        count_objects: Count objects within a polygonal or linear region based on their tracks.
+        display_counts: Display object counts on the frame.
+        process: Process input data and update counts.
+
+    Examples:
+        >>> counter = ObjectCounter()
+        >>> frame = cv2.imread("frame.jpg")
+        >>> results = counter.process(frame)
+        >>> print(f"Inward count: {counter.in_count}, Outward count: {counter.out_count}")
+    """
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Initialize the ObjectCounter class for real-time object counting in video streams."""
+        super().__init__(**kwargs)
+
+        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_count = defaultdict(lambda: {"IN": 0, "OUT": 0})  # Dictionary for counts, categorized by class
+        self.region_initialized = False  # Flag indicating whether the region has been initialized
+
+        self.show_in = self.CFG["show_in"]
+        self.show_out = self.CFG["show_out"]
+        self.margin = self.line_width * 2  # Scales the background rectangle size to display counts properly
+
+    def count_objects(
+        self,
+        current_centroid: tuple[float, float],
+        track_id: int,
+        prev_position: tuple[float, float] | None,
+        cls: int,
+    ) -> None:
+        """Count objects within a polygonal or linear region based on their tracks.
+
+        Args:
+            current_centroid (tuple[float, float]): Current centroid coordinates (x, y) in the current frame.
+            track_id (int): Unique identifier for the tracked object.
+            prev_position (tuple[float, float], optional): Last frame position coordinates (x, y) of the track.
+            cls (int): Class index for classwise count updates.
+
+        Examples:
+            >>> counter = ObjectCounter()
+            >>> track_line = {1: [100, 200], 2: [110, 210], 3: [120, 220]}
+            >>> box = [130, 230, 150, 250]
+            >>> track_id_num = 1
+            >>> previous_position = (120, 220)
+            >>> class_to_count = 0  # In COCO model, class 0 = person
+            >>> counter.count_objects((140, 240), track_id_num, previous_position, class_to_count)
+        """
+        if prev_position is None or track_id in self.counted_ids:
+            return
+
+        if len(self.region) == 2:  # Linear region (defined as a line segment)
+            if self.r_s.intersects(self.LineString([prev_position, current_centroid])):
+                # Determine orientation of the region (vertical or horizontal)
+                if abs(self.region[0][0] - self.region[1][0]) < abs(self.region[0][1] - self.region[1][1]):
+                    # Vertical region: Compare x-coordinates to determine direction
+                    if current_centroid[0] > prev_position[0]:  # Moving right
+                        self.in_count += 1
+                        self.classwise_count[self.names[cls]]["IN"] += 1
+                    else:  # Moving left
+                        self.out_count += 1
+                        self.classwise_count[self.names[cls]]["OUT"] += 1
+                # Horizontal region: Compare y-coordinates to determine direction
+                elif current_centroid[1] > prev_position[1]:  # Moving downward
+                    self.in_count += 1
+                    self.classwise_count[self.names[cls]]["IN"] += 1
+                else:  # Moving upward
+                    self.out_count += 1
+                    self.classwise_count[self.names[cls]]["OUT"] += 1
+                self.counted_ids.append(track_id)
+
+        elif len(self.region) > 2:  # Polygonal region
+            if self.r_s.contains(self.Point(current_centroid)):
+                # Determine motion direction for vertical or horizontal polygons
+                region_width = max(p[0] for p in self.region) - min(p[0] for p in self.region)
+                region_height = max(p[1] for p in self.region) - min(p[1] for p in self.region)
+
+                if (region_width < region_height and current_centroid[0] > prev_position[0]) or (
+                    region_width >= region_height and current_centroid[1] > prev_position[1]
+                ):  # Moving right or downward
+                    self.in_count += 1
+                    self.classwise_count[self.names[cls]]["IN"] += 1
+                else:  # Moving left or upward
+                    self.out_count += 1
+                    self.classwise_count[self.names[cls]]["OUT"] += 1
+                self.counted_ids.append(track_id)
+
+    def display_counts(self, plot_im) -> None:
+        """Display object counts on the input image or frame.
+
+        Args:
+            plot_im (np.ndarray): The image or frame to display counts on.
+
+        Examples:
+            >>> counter = ObjectCounter()
+            >>> frame = cv2.imread("image.jpg")
+            >>> counter.display_counts(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_count.items()
+            if (value["IN"] != 0 and self.show_in) or (value["OUT"] != 0 and self.show_out)
+        }
+        if labels_dict:
+            self.annotator.display_analytics(plot_im, labels_dict, (104, 31, 17), (255, 255, 255), self.margin)
+
+    def process(self, im0) -> SolutionResults:
+        """Process input data (frames or object tracks) and update object counts.
+
+        This method initializes the counting region, extracts tracks, draws bounding boxes and regions, updates object
+        counts, and displays the results on the input image.
+
+        Args:
+            im0 (np.ndarray): The input image or frame to be processed.
+
+        Returns:
+            (SolutionResults): Contains processed image `im0`, 'in_count' (int, count of objects entering the region),
+                'out_count' (int, count of objects exiting the region), 'classwise_count' (dict, per-class object
+                count), and 'total_tracks' (int, total number of tracked objects).
+
+        Examples:
+            >>> counter = ObjectCounter()
+            >>> frame = cv2.imread("path/to/image.jpg")
+            >>> results = counter.process(frame)
+        """
+        if not self.region_initialized:
+            self.initialize_region()
+            self.region_initialized = True
+
+        self.extract_tracks(im0)  # Extract tracks
+        self.annotator = SolutionAnnotator(im0, line_width=self.line_width)  # Initialize annotator
+
+        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, conf in zip(self.boxes, self.track_ids, self.clss, self.confs):
+            # Draw bounding box and counting region
+            self.annotator.box_label(box, label=self.adjust_box_label(cls, conf, track_id), color=colors(cls, True))
+            self.store_tracking_history(track_id, box)  # Store track history
+
+            # Store previous position of track for object counting
+            prev_position = None
+            if len(self.track_history[track_id]) > 1:
+                prev_position = self.track_history[track_id][-2]
+            self.count_objects(self.track_history[track_id][-1], track_id, prev_position, cls)  # object counting
+
+        plot_im = self.annotator.result()
+        self.display_counts(plot_im)  # Display the counts on the frame
+        self.display_output(plot_im)  # Display output with base class function
+
+        # Return SolutionResults
+        return SolutionResults(
+            plot_im=plot_im,
+            in_count=self.in_count,
+            out_count=self.out_count,
+            classwise_count=dict(self.classwise_count),
+            total_tracks=len(self.track_ids),
+        )

ultralytics/solutions/object_cropper.py

@@ -0,0 +1,87 @@
+# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
+
+from pathlib import Path
+from typing import Any
+
+from ultralytics.solutions.solutions import BaseSolution, SolutionResults
+from ultralytics.utils.plotting import save_one_box
+
+
+class ObjectCropper(BaseSolution):
+    """A class to manage the cropping of detected objects in a real-time video stream or images.
+
+    This class extends the BaseSolution class and provides functionality for cropping objects based on detected bounding
+    boxes. The cropped images are saved to a specified directory for further analysis or usage.
+
+    Attributes:
+        crop_dir (str): Directory where cropped object images are stored.
+        crop_idx (int): Counter for the total number of cropped objects.
+        iou (float): IoU (Intersection over Union) threshold for non-maximum suppression.
+        conf (float): Confidence threshold for filtering detections.
+
+    Methods:
+        process: Crop detected objects from the input image and save them to the output directory.
+
+    Examples:
+        >>> cropper = ObjectCropper()
+        >>> frame = cv2.imread("frame.jpg")
+        >>> processed_results = cropper.process(frame)
+        >>> print(f"Total cropped objects: {cropper.crop_idx}")
+    """
+
+    def __init__(self, **kwargs: Any) -> None:
+        """Initialize the ObjectCropper class for cropping objects from detected bounding boxes.
+
+        Args:
+            **kwargs (Any): Keyword arguments passed to the parent class and used for configuration including:
+                - crop_dir (str): Path to the directory for saving cropped object images.
+        """
+        super().__init__(**kwargs)
+
+        self.crop_dir = self.CFG["crop_dir"]  # Directory for storing cropped detections
+        Path(self.crop_dir).mkdir(parents=True, exist_ok=True)
+        if self.CFG["show"]:
+            self.LOGGER.warning(f"show=True is not supported for ObjectCropper; saving crops to '{self.crop_dir}'.")
+            self.CFG["show"] = False
+        self.crop_idx = 0  # Initialize counter for total cropped objects
+        self.iou = self.CFG["iou"]
+        self.conf = self.CFG["conf"]
+
+    def process(self, im0) -> SolutionResults:
+        """Crop detected objects from the input image and save them as separate images.
+
+        Args:
+            im0 (np.ndarray): The input image containing detected objects.
+
+        Returns:
+            (SolutionResults): A SolutionResults object containing the total number of cropped objects and processed
+                image.
+
+        Examples:
+            >>> cropper = ObjectCropper()
+            >>> frame = cv2.imread("image.jpg")
+            >>> results = cropper.process(frame)
+            >>> print(f"Total cropped objects: {results.total_crop_objects}")
+        """
+        with self.profilers[0]:
+            results = self.model.predict(
+                im0,
+                classes=self.classes,
+                conf=self.conf,
+                iou=self.iou,
+                device=self.CFG["device"],
+                verbose=False,
+            )[0]
+            self.clss = results.boxes.cls.tolist()  # required for logging only.
+
+        for box in results.boxes:
+            self.crop_idx += 1
+            save_one_box(
+                box.xyxy,
+                im0,
+                file=Path(self.crop_dir) / f"crop_{self.crop_idx}.jpg",
+                BGR=True,
+            )
+
+        # Return SolutionResults
+        return SolutionResults(plot_im=im0, total_crop_objects=self.crop_idx)