Skip to content

Register now for

# Reference for `ultralytics/solutions/distance_calculation.py`

Note

This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/solutions/distance_calculation.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!

## `` ultralytics.solutions.distance_calculation.DistanceCalculation

    DistanceCalculation(**kwargs: Any)

Bases: `BaseSolution`

A class to calculate distance between two objects in a real-time video stream based on their tracks.

This class extends BaseSolution to provide functionality for selecting objects and calculating the distance between them in a video stream using YOLO object detection and tracking.

Attributes:

Name | Type | Description
--- | --- | ---
`left_mouse_count` | `int` | Counter for left mouse button clicks.
`selected_boxes` | `dict[int, list[float]]` | Dictionary to store selected bounding boxes and their track IDs.
`centroids` | `list[list[int]]` | List to store centroids of selected bounding boxes.

Methods:

Name | Description
--- | ---
`mouse_event_for_distance` | Handle mouse events for selecting objects in the video stream.
`process` | Process video frames and calculate the distance between selected objects.

Examples:

    >>> distance_calc = DistanceCalculation()
    >>> frame = cv2.imread("frame.jpg")
    >>> results = distance_calc.process(frame)
    >>> cv2.imshow("Distance Calculation", results.plot_im)
    >>> cv2.waitKey(0)

Source code in `ultralytics/solutions/distance_calculation.py`

    36
    37
    38
    39
    40
    41
    42
    43
    def __init__(self, **kwargs: Any) -> None:
        """Initialize the DistanceCalculation class for measuring object distances in video streams."""
        super().__init__(**kwargs)

        # Mouse event information
        self.left_mouse_count = 0
        self.selected_boxes: dict[int, list[float]] = {}
        self.centroids: list[list[int]] = []  # Store centroids of selected objects

### `` mouse_event_for_distance

    mouse_event_for_distance(
        event: int, x: int, y: int, flags: int, param: Any
    ) -> None

Handle mouse events to select regions in a real-time video stream for distance calculation.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`event` | `int` | Type of mouse event (e.g., cv2.EVENT_MOUSEMOVE, cv2.EVENT_LBUTTONDOWN). | required
`x` | `int` | X-coordinate of the mouse pointer. | required
`y` | `int` | Y-coordinate of the mouse pointer. | required
`flags` | `int` | Flags associated with the event (e.g., cv2.EVENT_FLAG_CTRLKEY, cv2.EVENT_FLAG_SHIFTKEY). | required
`param` | `Any` | Additional parameters passed to the function. | required

Examples:

    >>> # Assuming 'dc' is an instance of DistanceCalculation
    >>> cv2.setMouseCallback("window_name", dc.mouse_event_for_distance)

Source code in `ultralytics/solutions/distance_calculation.py`

    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    def mouse_event_for_distance(self, event: int, x: int, y: int, flags: int, param: Any) -> None:
        """
        Handle mouse events to select regions in a real-time video stream for distance calculation.

        Args:
            event (int): Type of mouse event (e.g., cv2.EVENT_MOUSEMOVE, cv2.EVENT_LBUTTONDOWN).
            x (int): X-coordinate of the mouse pointer.
            y (int): Y-coordinate of the mouse pointer.
            flags (int): Flags associated with the event (e.g., cv2.EVENT_FLAG_CTRLKEY, cv2.EVENT_FLAG_SHIFTKEY).
            param (Any): Additional parameters passed to the function.

        Examples:
            >>> # Assuming 'dc' is an instance of DistanceCalculation
            >>> cv2.setMouseCallback("window_name", dc.mouse_event_for_distance)
        """
        if event == cv2.EVENT_LBUTTONDOWN:
            self.left_mouse_count += 1
            if self.left_mouse_count <= 2:
                for box, track_id in zip(self.boxes, self.track_ids):
                    if box[0] < x < box[2] and box[1] < y < box[3] and track_id not in self.selected_boxes:
                        self.selected_boxes[track_id] = box

        elif event == cv2.EVENT_RBUTTONDOWN:
            self.selected_boxes = {}
            self.left_mouse_count = 0

### `` process

    process(im0) -> SolutionResults

Process a video frame and calculate the distance between two selected bounding boxes.

This method extracts tracks from the input frame, annotates bounding boxes, and calculates the distance between two user-selected objects if they have been chosen.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`im0` | `ndarray` | The input image frame to process. | required

Returns:

Type | Description
--- | ---
`SolutionResults` | Contains processed image `plot_im`, `total_tracks` (int) representing the total number of tracked objects, and `pixels_distance` (float) representing the distance between selected objects in pixels.

Examples:

    >>> import numpy as np
    >>> from ultralytics.solutions import DistanceCalculation
    >>> dc = DistanceCalculation()
    >>> frame = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
    >>> results = dc.process(frame)
    >>> print(f"Distance: {results.pixels_distance:.2f} pixels")

Source code in `ultralytics/solutions/distance_calculation.py`

     71
     72
     73
     74
     75
     76
     77
     78
     79
     80
     81
     82
     83
     84
     85
     86
     87
     88
     89
     90
     91
     92
     93
     94
     95
     96
     97
     98
     99
    100
    101
    102
    103
    104
    105
    106
    107
    108
    109
    110
    111
    112
    113
    114
    115
    116
    117
    118
    119
    120
    121
    122
    123
    124
    125
    126
    def process(self, im0) -> SolutionResults:
        """
        Process a video frame and calculate the distance between two selected bounding boxes.

        This method extracts tracks from the input frame, annotates bounding boxes, and calculates the distance
        between two user-selected objects if they have been chosen.

        Args:
            im0 (np.ndarray): The input image frame to process.

        Returns:
            (SolutionResults): Contains processed image `plot_im`, `total_tracks` (int) representing the total number
                of tracked objects, and `pixels_distance` (float) representing the distance between selected objects
                in pixels.

        Examples:
            >>> import numpy as np
            >>> from ultralytics.solutions import DistanceCalculation
            >>> dc = DistanceCalculation()
            >>> frame = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
            >>> results = dc.process(frame)
            >>> print(f"Distance: {results.pixels_distance:.2f} pixels")
        """
        self.extract_tracks(im0)  # Extract tracks
        annotator = SolutionAnnotator(im0, line_width=self.line_width)  # Initialize annotator

        pixels_distance = 0
        # 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):
            annotator.box_label(box, color=colors(int(cls), True), label=self.adjust_box_label(cls, conf, track_id))

            # Update selected boxes if they're being tracked
            if len(self.selected_boxes) == 2:
                for trk_id in self.selected_boxes.keys():
                    if trk_id == track_id:
                        self.selected_boxes[track_id] = box

        if len(self.selected_boxes) == 2:
            # Calculate centroids of selected boxes
            self.centroids.extend(
                [[int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2)] for box in self.selected_boxes.values()]
            )
            # Calculate Euclidean distance between centroids
            pixels_distance = math.sqrt(
                (self.centroids[0][0] - self.centroids[1][0]) ** 2 + (self.centroids[0][1] - self.centroids[1][1]) ** 2
            )
            annotator.plot_distance_and_line(pixels_distance, self.centroids)

        self.centroids = []  # Reset centroids for next frame
        plot_im = annotator.result()
        self.display_output(plot_im)  # Display output with base class function
        if self.CFG.get("show") and self.env_check:
            cv2.setMouseCallback("Ultralytics Solutions", self.mouse_event_for_distance)

        # Return SolutionResults with processed image and calculated metrics
        return SolutionResults(plot_im=plot_im, pixels_distance=pixels_distance, total_tracks=len(self.track_ids))

📅 Created 1 year ago ✏️ Updated 1 year ago