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

Refactor all Ultralytics Solutions (#12790)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: RizwanMunawar <chr043416@gmail.com>
This commit is contained in:
Glenn Jocher 2024-05-18 18:14:42 +02:00 • committed by GitHub
parent a2ecb24176
commit 2af71d15a6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
134 changed files with 845 additions and 1020 deletions
Download patch file Download diff file Expand all files Collapse all files

139
ultralytics/solutions/distance_calculation.py
View file

@ -9,39 +9,9 @@ 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."""
"""A class to calculate distance between two objects in a 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(
def __init__(
self,
names,
pixels_per_meter=10,
@ -51,52 +21,66 @@ class DistanceCalculation:
centroid_color=(255, 0, 255),
):
"""
Configures the distance calculation and display parameters.
Initializes the DistanceCalculation class with the given 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
names (dict): Dictionary mapping class indices to class 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).
centroid_color (tuple, optional): Color of the centroids drawn (BGR format). Defaults to (255, 0, 255).
"""
self.names = names
self.pixel_per_meter = pixels_per_meter
# Visual & image information
self.im0 = None
self.annotator = None
self.view_img = view_img
self.line_thickness = line_thickness
self.line_color = line_color
self.centroid_color = centroid_color
# Prediction & tracking information
self.clss = None
self.names = names
self.boxes = None
self.line_thickness = line_thickness
self.trk_ids = None
# Distance calculation information
self.centroids = []
self.pixel_per_meter = pixels_per_meter
# Mouse event information
self.left_mouse_count = 0
self.selected_boxes = {}
# Check if environment supports imshow
self.env_check = check_imshow(warn=True)
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.
Handles mouse events to select regions 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.
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.
"""
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:
elif event == cv2.EVENT_RBUTTONDOWN:
self.selected_boxes = {}
self.left_mouse_count = 0
def extract_tracks(self, tracks):
"""
Extracts results from the provided data.
Extracts tracking results from the provided data.
Args:
tracks (list): List of tracks obtained from the object tracking process.
@ -105,55 +89,65 @@ class DistanceCalculation:
self.clss = tracks[0].boxes.cls.cpu().tolist()
self.trk_ids = tracks[0].boxes.id.int().cpu().tolist()
def calculate_centroid(self, box):
@staticmethod
def calculate_centroid(box):
"""
Calculate the centroid of bounding box.
Calculates the centroid of a bounding box.
Args:
box (list): Bounding box data
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):
"""
Calculate distance between two centroids.
Calculates the distance between two centroids.
Args:
centroid1 (point): First bounding box data
centroid2 (point): Second bounding box data
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)
return pixel_distance / self.pixel_per_meter, (pixel_distance / self.pixel_per_meter) * 1000
distance_m = pixel_distance / self.pixel_per_meter
distance_mm = distance_m * 1000
return distance_m, distance_mm
def start_process(self, im0, tracks):
"""
Calculate distance between two bounding boxes based on tracking data.
Processes the video frame and calculates the distance between two bounding boxes.
Args:
im0 (nd array): Image
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
self.extract_tracks(tracks)
return im0
self.annotator = Annotator(self.im0, line_width=2)
self.extract_tracks(tracks)
self.annotator = Annotator(self.im0, line_width=self.line_thickness)
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():
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:
for trk_id, box in self.selected_boxes.items():
centroid = self.calculate_centroid(self.selected_boxes[trk_id])
self.centroids.append(centroid)
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(
@ -168,7 +162,7 @@ class DistanceCalculation:
return im0
def display_frames(self):
"""Display frame."""
"""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)
@ -178,4 +172,5 @@ class DistanceCalculation:
if __name__ == "__main__":
DistanceCalculation()
names = {0: "person", 1: "car"} # example class names
distance_calculation = DistanceCalculation(names)