HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.2.88 Update distance-calculation to pixels (#15984)

Browse source 
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Muhammad Rizwan Munawar 2024-09-06 02:00:55 +05:00 • committed by GitHub
parent 41dfd65cc1
commit 4673fae31d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 43 additions and 87 deletions
Download patch file Download diff file Expand all files Collapse all files

75
ultralytics/solutions/distance_calculation.py
View file

@ -14,18 +14,16 @@ class DistanceCalculation:
def __init__(
self,
names,
pixels_per_meter=10,
view_img=False,
line_thickness=2,
line_color=(255, 255, 0),
centroid_color=(255, 0, 255),
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.
pixels_per_meter (int, optional): Conversion factor from pixels to meters. Defaults to 10.
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).
@ -39,7 +37,6 @@ class DistanceCalculation:
self.centroid_color = centroid_color
# Prediction & tracking information
self.clss = None
self.names = names
self.boxes = None
self.line_thickness = line_thickness
@ -47,7 +44,6 @@ class DistanceCalculation:
# Distance calculation information
self.centroids = []
self.pixel_per_meter = pixels_per_meter
# Mouse event information
self.left_mouse_count = 0
@ -55,6 +51,7 @@ class DistanceCalculation:
# Check if environment supports imshow
self.env_check = check_imshow(warn=True)
self.window_name = "Ultralytics Solutions"
def mouse_event_for_distance(self, event, x, y, flags, param):
"""
@ -78,46 +75,6 @@ class DistanceCalculation:
self.selected_boxes = {}
self.left_mouse_count = 0
def extract_tracks(self, tracks):
"""
Extracts tracking 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()
@staticmethod
def calculate_centroid(box):
"""
Calculates the centroid of a bounding box.
Args:
box (list): Bounding box coordinates [x1, y1, x2, y2].
Returns:
(tuple): Centroid coordinates (x, y).
"""
return int((box[0] + box[2]) // 2), int((box[1] + box[3]) // 2)
def calculate_distance(self, centroid1, centroid2):
"""
Calculates the distance between two centroids.
Args:
centroid1 (tuple): Coordinates of the first centroid (x, y).
centroid2 (tuple): Coordinates of the second centroid (x, y).
Returns:
(tuple): Distance in meters and millimeters.
"""
pixel_distance = math.sqrt((centroid1[0] - centroid2[0]) ** 2 + (centroid1[1] - centroid2[1]) ** 2)
distance_m = pixel_distance / self.pixel_per_meter
distance_mm = distance_m * 1000
return distance_m, distance_mm
def start_process(self, im0, tracks):
"""
Processes the video frame and calculates the distance between two bounding boxes.
@ -135,10 +92,13 @@ class DistanceCalculation:
self.display_frames()
return im0
self.extract_tracks(tracks)
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, self.clss, self.trk_ids):
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:
@ -147,12 +107,15 @@ class DistanceCalculation:
self.selected_boxes[track_id] = box
if len(self.selected_boxes) == 2:
self.centroids = [self.calculate_centroid(self.selected_boxes[trk_id]) for trk_id in self.selected_boxes]
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
# 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 = []
@ -163,9 +126,9 @@ class DistanceCalculation:
def display_frames(self):
"""Displays the current frame with annotations."""
cv2.namedWindow("Ultralytics Distance Estimation")
cv2.setMouseCallback("Ultralytics Distance Estimation", self.mouse_event_for_distance)
cv2.imshow("Ultralytics Distance Estimation", self.im0)
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