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

Update distance-calculation solution (#16907)

Browse source 
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
Muhammad Rizwan Munawar 2024-10-14 19:44:36 +05:00 • committed by GitHub
parent 1052cf41f8
commit 8e3846d377
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 41 additions and 106 deletions
Download patch file Download diff file Expand all files Collapse all files

87
ultralytics/solutions/distance_calculation.py
View file

@ -4,55 +4,21 @@ import math
import cv2
from ultralytics.utils.checks import check_imshow
from ultralytics.solutions.solutions import BaseSolution # Import a parent class
from ultralytics.utils.plotting import Annotator, colors
class DistanceCalculation:
class DistanceCalculation(BaseSolution):
"""A class to calculate distance between two objects in a real-time video stream based on their tracks."""
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 = []
def __init__(self, **kwargs):
"""Initializes the DistanceCalculation class with the given parameters."""
super().__init__(**kwargs)
# 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"
def mouse_event_for_distance(self, event, x, y, flags, param):
"""
Handles mouse events to select regions in a real-time video stream.
@ -67,7 +33,7 @@ class DistanceCalculation:
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):
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
@ -75,30 +41,21 @@ class DistanceCalculation:
self.selected_boxes = {}
self.left_mouse_count = 0
def start_process(self, im0, tracks):
def calculate(self, im0):
"""
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.annotator = Annotator(im0, line_width=self.line_width) # Initialize annotator
self.extract_tracks(im0) # Extract 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, clss, self.trk_ids):
# 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:
@ -115,25 +72,11 @@ class DistanceCalculation:
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.annotator.plot_distance_and_line(pixels_distance, self.centroids)
self.centroids = []
if self.view_img and self.env_check:
self.display_frames()
self.display_output(im0) # display output with base class function
cv2.setMouseCallback("Ultralytics Solutions", self.mouse_event_for_distance)
return im0
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
if __name__ == "__main__":
names = {0: "person", 1: "car"} # example class names
distance_calculation = DistanceCalculation(names)
return im0 # return output image for more usage

4
ultralytics/solutions/object_counter.py
View file

@ -112,13 +112,13 @@ class ObjectCounter(BaseSolution):
# Iterate over bounding boxes, track ids and classes index
for box, track_id, cls in zip(self.boxes, self.track_ids, self.clss):
# Draw bounding box and counting region
self.annotator.box_label(box, label=self.names[cls], color=colors(track_id, True))
self.annotator.box_label(box, label=self.names[cls], color=colors(cls, True))
self.store_tracking_history(track_id, box) # Store track history
self.store_classwise_counts(cls) # store classwise counts in dict
# Draw tracks of objects
self.annotator.draw_centroid_and_tracks(
self.track_line, color=colors(int(track_id), True), track_thickness=self.line_width
self.track_line, color=colors(int(cls), True), track_thickness=self.line_width
)
# store previous position of track for object counting