/
distance_calculation.py
181 lines (145 loc) 路 6.19 KB
/
distance_calculation.py
1
2
3
4
5
6
7
8
9
10
11
12
13
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
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
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
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
# Ultralytics YOLO 馃殌, AGPL-3.0 license
import math
import cv2
from ultralytics.utils.checks import check_imshow
from ultralytics.utils.plotting import Annotator, colors
class DistanceCalculation:
"""A class to calculate distance between two objects in real-time video stream based on their tracks."""
def __init__(self):
"""Initializes the distance calculation class with default values for Visual, Image, track and distance
parameters.
"""
# Visual & im0 information
self.im0 = None
self.annotator = None
self.view_img = False
self.line_color = (255, 255, 0)
self.centroid_color = (255, 0, 255)
# Predict/track information
self.clss = None
self.names = None
self.boxes = None
self.line_thickness = 2
self.trk_ids = None
# Distance calculation information
self.centroids = []
self.pixel_per_meter = 10
# Mouse event
self.left_mouse_count = 0
self.selected_boxes = {}
# Check if environment support imshow
self.env_check = check_imshow(warn=True)
def set_args(
self,
names,
pixels_per_meter=10,
view_img=False,
line_thickness=2,
line_color=(255, 255, 0),
centroid_color=(255, 0, 255),
):
"""
Configures the distance calculation and display parameters.
Args:
names (dict): object detection classes names
pixels_per_meter (int): Number of pixels in meter
view_img (bool): Flag indicating frame display
line_thickness (int): Line thickness for bounding boxes.
line_color (RGB): color of centroids line
centroid_color (RGB): colors of bbox centroids
"""
self.names = names
self.pixel_per_meter = pixels_per_meter
self.view_img = view_img
self.line_thickness = line_thickness
self.line_color = line_color
self.centroid_color = centroid_color
def mouse_event_for_distance(self, event, x, y, flags, param):
"""
This function is designed to move region with mouse events in a real-time video stream.
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 flags associated with the event (e.g., cv2.EVENT_FLAG_CTRLKEY,
cv2.EVENT_FLAG_SHIFTKEY, etc.).
param (dict): Additional parameters you may want to pass to the function.
"""
global selected_boxes
global left_mouse_count
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] = []
self.selected_boxes[track_id] = box
if event == cv2.EVENT_RBUTTONDOWN:
self.selected_boxes = {}
self.left_mouse_count = 0
def extract_tracks(self, tracks):
"""
Extracts results from the provided data.
Args:
tracks (list): List of tracks obtained from the object tracking process.
"""
self.boxes = tracks[0].boxes.xyxy.cpu()
self.clss = tracks[0].boxes.cls.cpu().tolist()
self.trk_ids = tracks[0].boxes.id.int().cpu().tolist()
def calculate_centroid(self, box):
"""
Calculate the centroid of bounding box.
Args:
box (list): Bounding box data
"""
return int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2)
def calculate_distance(self, centroid1, centroid2):
"""
Calculate distance between two centroids.
Args:
centroid1 (point): First bounding box data
centroid2 (point): Second bounding box data
"""
pixel_distance = math.sqrt((centroid1[0] - centroid2[0]) ** 2 + (centroid1[1] - centroid2[1]) ** 2)
return pixel_distance / self.pixel_per_meter, (pixel_distance / self.pixel_per_meter) * 1000
def start_process(self, im0, tracks):
"""
Calculate distance between two bounding boxes based on tracking data.
Args:
im0 (nd array): Image
tracks (list): List of tracks obtained from the object tracking process.
"""
self.im0 = im0
if tracks[0].boxes.id is None:
if self.view_img:
self.display_frames()
return
self.extract_tracks(tracks)
self.annotator = Annotator(self.im0, line_width=2)
for box, cls, track_id in zip(self.boxes, self.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.items():
if trk_id == track_id:
self.selected_boxes[track_id] = box
if len(self.selected_boxes) == 2:
for trk_id, box in self.selected_boxes.items():
centroid = self.calculate_centroid(self.selected_boxes[trk_id])
self.centroids.append(centroid)
distance_m, distance_mm = self.calculate_distance(self.centroids[0], self.centroids[1])
self.annotator.plot_distance_and_line(
distance_m, distance_mm, self.centroids, self.line_color, self.centroid_color
)
self.centroids = []
if self.view_img and self.env_check:
self.display_frames()
return im0
def display_frames(self):
"""Display frame."""
cv2.namedWindow("Ultralytics Distance Estimation")
cv2.setMouseCallback("Ultralytics Distance Estimation", self.mouse_event_for_distance)
cv2.imshow("Ultralytics Distance Estimation", self.im0)
if cv2.waitKey(1) & 0xFF == ord("q"):
return
if __name__ == "__main__":
DistanceCalculation()