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(names, view_img=False, line_thickness=2, line_color=(255, 0, 255), centroid_color=(104, 31, 17))

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

Parameters:

Name Type Description Default
names dict

Dictionary of classes names.

 required
view_img bool

Flag to indicate if the video stream should be displayed. Defaults to False.

 False
line_thickness int

Thickness of the lines drawn on the image. Defaults to 2.

 2
line_color tuple

Color of the lines drawn on the image (BGR format). Defaults to (255, 255, 0).

 (255, 0, 255)
centroid_color tuple

Color of the centroids drawn (BGR format). Defaults to (255, 0, 255).

 (104, 31, 17)
Source code in ultralytics/solutions/distance_calculation.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
def __init__(
    self,
    names,
    view_img=False,
    line_thickness=2,
    line_color=(255, 0, 255),
    centroid_color=(104, 31, 17),
):
    """
    Initializes the DistanceCalculation class with the given parameters.

    Args:
        names (dict): Dictionary of classes names.
        view_img (bool, optional): Flag to indicate if the video stream should be displayed. Defaults to False.
        line_thickness (int, optional): Thickness of the lines drawn on the image. Defaults to 2.
        line_color (tuple, optional): Color of the lines drawn on the image (BGR format). Defaults to (255, 255, 0).
        centroid_color (tuple, optional): Color of the centroids drawn (BGR format). Defaults to (255, 0, 255).
    """
    # Visual & image information
    self.im0 = None
    self.annotator = None
    self.view_img = view_img
    self.line_color = line_color
    self.centroid_color = centroid_color

    # Prediction & tracking information
    self.names = names
    self.boxes = None
    self.line_thickness = line_thickness
    self.trk_ids = None

    # Distance calculation information
    self.centroids = []

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

    # Check if environment supports imshow
    self.env_check = check_imshow(warn=True)
    self.window_name = "Ultralytics Solutions"

display_frames 

display_frames()

Displays the current frame with annotations.

Source code in ultralytics/solutions/distance_calculation.py 
127
128
129
130
131
132
133
134
def display_frames(self):
    """Displays the current frame with annotations."""
    cv2.namedWindow(self.window_name)
    cv2.setMouseCallback(self.window_name, self.mouse_event_for_distance)
    cv2.imshow(self.window_name, self.im0)

    if cv2.waitKey(1) & 0xFF == ord("q"):
        return

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.

Parameters:

Name Type Description Default
event int

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

 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, etc.).

 required
param dict

Additional parameters passed to the function.

 required
Source code in ultralytics/solutions/distance_calculation.py 
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
def mouse_event_for_distance(self, event, x, y, flags, param):
    """
    Handles mouse events to select regions in a real-time video stream.

    Args:
        event (int): Type of mouse event (e.g., cv2.EVENT_MOUSEMOVE, cv2.EVENT_LBUTTONDOWN, etc.).
        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, etc.).
        param (dict): Additional parameters passed to the function.
    """
    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.trk_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

start_process 

start_process(im0, tracks)

Processes the video frame and calculates the distance between two bounding boxes.

Parameters:

Name Type Description Default
im0 ndarray

The image frame.

 required
tracks list

List of tracks obtained from the object tracking process.

 required

Returns:

Type Description
ndarray

The processed image frame.
Source code in ultralytics/solutions/distance_calculation.py 
 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
def start_process(self, im0, tracks):
    """
    Processes the video frame and calculates the distance between two bounding boxes.

    Args:
        im0 (ndarray): The image frame.
        tracks (list): List of tracks obtained from the object tracking process.

    Returns:
        (ndarray): The processed image frame.
    """
    self.im0 = im0
    if tracks[0].boxes.id is None:
        if self.view_img:
            self.display_frames()
        return im0

    self.boxes = tracks[0].boxes.xyxy.cpu()
    clss = tracks[0].boxes.cls.cpu().tolist()
    self.trk_ids = tracks[0].boxes.id.int().cpu().tolist()

    self.annotator = Annotator(self.im0, line_width=self.line_thickness)

    for box, cls, track_id in zip(self.boxes, clss, self.trk_ids):
        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.line_color, self.centroid_color)

    self.centroids = []

    if self.view_img and self.env_check:
        self.display_frames()

    return im0




📅 Created 9 months ago ✏️ Updated 1 month ago
jk4e glenn-jocher Burhan-Q AyushExel RizwanMunawar