Skip to content

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)

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:

NameTypeDescription
left_mouse_countint

Counter for left mouse button clicks.

selected_boxesDict[int, List[float]]

Dictionary to store selected bounding boxes and their track IDs.

annotatorAnnotator

An instance of the Annotator class for drawing on the image.

boxesList[List[float]]

List of bounding boxes for detected objects.

track_idsList[int]

List of track IDs for detected objects.

clssList[int]

List of class indices for detected objects.

namesList[str]

List of class names that the model can detect.

centroidsList[List[int]]

List to store centroids of selected bounding boxes.

Methods:

NameDescription
mouse_event_for_distance

Handles mouse events for selecting objects in the video stream.

calculate

Processes video frames and calculates the distance between selected objects.

Examples:

>>> distance_calc = DistanceCalculation()
>>> frame = cv2.imread("frame.jpg")
>>> processed_frame = distance_calc.calculate(frame)
>>> cv2.imshow("Distance Calculation", processed_frame)
>>> cv2.waitKey(0)
Source code in ultralytics/solutions/distance_calculation.py 
40
41
42
43
44
45
46
def __init__(self, **kwargs):
    """Initializes the DistanceCalculation class for measuring object distances in video streams."""
    super().__init__(**kwargs)

    # Mouse event information
    self.left_mouse_count = 0
    self.selected_boxes = {}

calculate

calculate(im0)

Processes a video frame and calculates 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:

NameTypeDescriptionDefault
im0ndarray

The input image frame to process.

required

Returns:

TypeDescription
ndarray

The processed image frame with annotations and distance calculations.

Examples:

>>> import numpy as np
>>> from ultralytics.solutions import DistanceCalculation
>>> dc = DistanceCalculation()
>>> frame = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
>>> processed_frame = dc.calculate(frame)
Source code in ultralytics/solutions/distance_calculation.py 
 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
def calculate(self, im0):
    """
    Processes a video frame and calculates 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 (numpy.ndarray): The input image frame to process.

    Returns:
        (numpy.ndarray): The processed image frame with annotations and distance calculations.

    Examples:
        >>> import numpy as np
        >>> from ultralytics.solutions import DistanceCalculation
        >>> dc = DistanceCalculation()
        >>> frame = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
        >>> processed_frame = dc.calculate(frame)
    """
    self.annotator = Annotator(im0, line_width=self.line_width)  # Initialize annotator
    self.extract_tracks(im0)  # Extract tracks

    # Iterate over bounding boxes, track ids and classes index
    for box, track_id, cls in zip(self.boxes, self.track_ids, self.clss):
        self.annotator.box_label(box, color=colors(int(cls), True), label=self.names[int(cls)])

        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:
        # Store user selected boxes in centroids list
        self.centroids.extend(
            [[int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2)] for box in self.selected_boxes.values()]
        )
        # Calculate pixels distance
        pixels_distance = math.sqrt(
            (self.centroids[0][0] - self.centroids[1][0]) ** 2 + (self.centroids[0][1] - self.centroids[1][1]) ** 2
        )
        self.annotator.plot_distance_and_line(pixels_distance, self.centroids)

    self.centroids = []

    self.display_output(im0)  # display output with base class function
    cv2.setMouseCallback("Ultralytics Solutions", self.mouse_event_for_distance)

    return im0  # return output image for more usage

mouse_event_for_distance

mouse_event_for_distance(event, x, y, flags, param)

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

Parameters:

NameTypeDescriptionDefault
eventint

Type of mouse event (e.g., cv2.EVENT_MOUSEMOVE, cv2.EVENT_LBUTTONDOWN).

required
xint

X-coordinate of the mouse pointer.

required
yint

Y-coordinate of the mouse pointer.

required
flagsint

Flags associated with the event (e.g., cv2.EVENT_FLAG_CTRLKEY, cv2.EVENT_FLAG_SHIFTKEY).

required
paramDict

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 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
def mouse_event_for_distance(self, event, x, y, flags, param):
    """
    Handles 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 (Dict): 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



📅 Created 10 months ago ✏️ Updated 2 months ago
jk4e glenn-jocher Burhan-Q AyushExel RizwanMunawar