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

18
ultralytics/solutions/__init__.py
View file

@ -1 +1,19 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license
from .ai_gym import AIGym
from .distance_calculation import DistanceCalculation
from .heatmap import Heatmap
from .object_counter import ObjectCounter
from .parking_management import ParkingManagement
from .queue_management import QueueManager
from .speed_estimation import SpeedEstimator
__all__ = (
"AIGym",
"DistanceCalculation",
"Heatmap",
"ObjectCounter",
"ParkingManagement",
"QueueManager",
"SpeedEstimator",
)

136
ultralytics/solutions/ai_gym.py
View file

@ -9,34 +9,7 @@ from ultralytics.utils.plotting import Annotator
class AIGym:
"""A class to manage the gym steps of people in a real-time video stream based on their poses."""
def __init__(self):
"""Initializes the AIGym with default values for Visual and Image parameters."""
# Image and line thickness
self.im0 = None
self.tf = None
# Keypoints and count information
self.keypoints = None
self.poseup_angle = None
self.posedown_angle = None
self.threshold = 0.001
# Store stage, count and angle information
self.angle = None
self.count = None
self.stage = None
self.pose_type = "pushup"
self.kpts_to_check = None
# Visual Information
self.view_img = False
self.annotator = None
# Check if environment support imshow
self.env_check = check_imshow(warn=True)
def set_args(
def __init__(
self,
kpts_to_check,
line_thickness=2,
@ -46,22 +19,40 @@ class AIGym:
pose_type="pullup",
):
"""
Configures the AIGym line_thickness, save image and view image parameters.
Initializes the AIGym class with the specified parameters.
Args:
kpts_to_check (list): 3 keypoints for counting
line_thickness (int): Line thickness for bounding boxes.
view_img (bool): display the im0
pose_up_angle (float): Angle to set pose position up
pose_down_angle (float): Angle to set pose position down
pose_type (str): "pushup", "pullup" or "abworkout"
kpts_to_check (list): Indices of keypoints to check.
line_thickness (int, optional): Thickness of the lines drawn. Defaults to 2.
view_img (bool, optional): Flag to display the image. Defaults to False.
pose_up_angle (float, optional): Angle threshold for the 'up' pose. Defaults to 145.0.
pose_down_angle (float, optional): Angle threshold for the 'down' pose. Defaults to 90.0.
pose_type (str, optional): Type of pose to detect ('pullup', 'pushup', 'abworkout'). Defaults to "pullup".
"""
self.kpts_to_check = kpts_to_check
# Image and line thickness
self.im0 = None
self.tf = line_thickness
self.view_img = view_img
# Keypoints and count information
self.keypoints = None
self.poseup_angle = pose_up_angle
self.posedown_angle = pose_down_angle
self.threshold = 0.001
# Store stage, count and angle information
self.angle = None
self.count = None
self.stage = None
self.pose_type = pose_type
self.kpts_to_check = kpts_to_check
# Visual Information
self.view_img = view_img
self.annotator = None
# Check if environment supports imshow
self.env_check = check_imshow(warn=True)
def start_counting(self, im0, results, frame_count):
"""
@ -69,19 +60,24 @@ class AIGym:
Args:
im0 (ndarray): Current frame from the video stream.
results (list): Pose estimation data
frame_count (int): store current frame count
results (list): Pose estimation data.
frame_count (int): Current frame count.
"""
self.im0 = im0
# Initialize count, angle, and stage lists on the first frame
if frame_count == 1:
self.count = [0] * len(results[0])
self.angle = [0] * len(results[0])
self.stage = ["-" for _ in results[0]]
self.keypoints = results[0].keypoints.data
self.annotator = Annotator(im0, line_width=2)
for ind, k in enumerate(reversed(self.keypoints)):
if self.pose_type in {"pushup", "pullup"}:
# Estimate angle and draw specific points based on pose type
if self.pose_type in {"pushup", "pullup", "abworkout"}:
self.angle[ind] = self.annotator.estimate_pose_angle(
k[int(self.kpts_to_check[0])].cpu(),
k[int(self.kpts_to_check[1])].cpu(),
@ -89,55 +85,32 @@ class AIGym:
)
self.im0 = self.annotator.draw_specific_points(k, self.kpts_to_check, shape=(640, 640), radius=10)
if self.pose_type == "abworkout":
self.angle[ind] = self.annotator.estimate_pose_angle(
k[int(self.kpts_to_check[0])].cpu(),
k[int(self.kpts_to_check[1])].cpu(),
k[int(self.kpts_to_check[2])].cpu(),
)
self.im0 = self.annotator.draw_specific_points(k, self.kpts_to_check, shape=(640, 640), radius=10)
if self.angle[ind] > self.poseup_angle:
self.stage[ind] = "down"
if self.angle[ind] < self.posedown_angle and self.stage[ind] == "down":
self.stage[ind] = "up"
self.count[ind] += 1
# Check and update pose stages and counts based on angle
if self.pose_type in {"abworkout", "pullup"}:
if self.angle[ind] > self.poseup_angle:
self.stage[ind] = "down"
if self.angle[ind] < self.posedown_angle and self.stage[ind] == "down":
self.stage[ind] = "up"
self.count[ind] += 1
elif self.pose_type == "pushup":
if self.angle[ind] > self.poseup_angle:
self.stage[ind] = "up"
if self.angle[ind] < self.posedown_angle and self.stage[ind] == "up":
self.stage[ind] = "down"
self.count[ind] += 1
self.annotator.plot_angle_and_count_and_stage(
angle_text=self.angle[ind],
count_text=self.count[ind],
stage_text=self.stage[ind],
center_kpt=k[int(self.kpts_to_check[1])],
line_thickness=self.tf,
)
if self.pose_type == "pushup":
if self.angle[ind] > self.poseup_angle:
self.stage[ind] = "up"
if self.angle[ind] < self.posedown_angle and self.stage[ind] == "up":
self.stage[ind] = "down"
self.count[ind] += 1
self.annotator.plot_angle_and_count_and_stage(
angle_text=self.angle[ind],
count_text=self.count[ind],
stage_text=self.stage[ind],
center_kpt=k[int(self.kpts_to_check[1])],
line_thickness=self.tf,
)
if self.pose_type == "pullup":
if self.angle[ind] > self.poseup_angle:
self.stage[ind] = "down"
if self.angle[ind] < self.posedown_angle and self.stage[ind] == "down":
self.stage[ind] = "up"
self.count[ind] += 1
self.annotator.plot_angle_and_count_and_stage(
angle_text=self.angle[ind],
count_text=self.count[ind],
stage_text=self.stage[ind],
center_kpt=k[int(self.kpts_to_check[1])],
line_thickness=self.tf,
)
# Draw keypoints
self.annotator.kpts(k, shape=(640, 640), radius=1, kpt_line=True)
# Display the image if environment supports it and view_img is True
if self.env_check and self.view_img:
cv2.imshow("Ultralytics YOLOv8 AI GYM", self.im0)
if cv2.waitKey(1) & 0xFF == ord("q"):
@ -147,4 +120,5 @@ class AIGym:
if __name__ == "__main__":
AIGym()
kpts_to_check = [0, 1, 2] # example keypoints
aigym = AIGym(kpts_to_check)

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)

244
ultralytics/solutions/heatmap.py
View file

@ -16,62 +16,11 @@ from shapely.geometry import LineString, Point, Polygon
class Heatmap:
"""A class to draw heatmaps in real-time video stream based on their tracks."""
def __init__(self):
"""Initializes the heatmap class with default values for Visual, Image, track, count and heatmap parameters."""
# Visual information
self.annotator = None
self.view_img = False
self.shape = "circle"
self.names = None # Classes names
# Image information
self.imw = None
self.imh = None
self.im0 = None
self.tf = 2
self.view_in_counts = True
self.view_out_counts = True
# Heatmap colormap and heatmap np array
self.colormap = None
self.heatmap = None
self.heatmap_alpha = 0.5
# Predict/track information
self.boxes = None
self.track_ids = None
self.clss = None
self.track_history = defaultdict(list)
# Region & Line Information
self.count_reg_pts = None
self.counting_region = None
self.line_dist_thresh = 15
self.region_thickness = 5
self.region_color = (255, 0, 255)
# Object Counting Information
self.in_counts = 0
self.out_counts = 0
self.count_ids = []
self.class_wise_count = {}
self.count_txt_color = (0, 0, 0)
self.count_bg_color = (255, 255, 255)
self.cls_txtdisplay_gap = 50
# Decay factor
self.decay_factor = 0.99
# Check if environment support imshow
self.env_check = check_imshow(warn=True)
def set_args(
def __init__(
self,
imw,
imh,
classes_names=None,
classes_names,
imw=0,
imh=0,
colormap=cv2.COLORMAP_JET,
heatmap_alpha=0.5,
view_img=False,
@ -87,71 +36,78 @@ class Heatmap:
decay_factor=0.99,
shape="circle",
):
"""
Configures the heatmap colormap, width, height and display parameters.
"""Initializes the heatmap class with default values for Visual, Image, track, count and heatmap parameters."""
Args:
colormap (cv2.COLORMAP): The colormap to be set.
imw (int): The width of the frame.
imh (int): The height of the frame.
classes_names (dict): Classes names
line_thickness (int): Line thickness for bounding boxes.
heatmap_alpha (float): alpha value for heatmap display
view_img (bool): Flag indicating frame display
view_in_counts (bool): Flag to control whether to display the incounts on video stream.
view_out_counts (bool): Flag to control whether to display the outcounts on video stream.
count_reg_pts (list): Object counting region points
count_txt_color (RGB color): count text color value
count_bg_color (RGB color): count highlighter line color
count_reg_color (RGB color): Color of object counting region
region_thickness (int): Object counting Region thickness
line_dist_thresh (int): Euclidean Distance threshold for line counter
decay_factor (float): value for removing heatmap area after object passed
shape (str): Heatmap shape, rect or circle shape supported
"""
self.tf = line_thickness
self.names = classes_names
# Visual information
self.annotator = None
self.view_img = view_img
self.shape = shape
self.initialized = False
self.names = classes_names # Classes names
# Image information
self.imw = imw
self.imh = imh
self.heatmap_alpha = heatmap_alpha
self.view_img = view_img
self.im0 = None
self.tf = line_thickness
self.view_in_counts = view_in_counts
self.view_out_counts = view_out_counts
# Heatmap colormap and heatmap np array
self.colormap = colormap
self.heatmap = None
self.heatmap_alpha = heatmap_alpha
# Predict/track information
self.boxes = None
self.track_ids = None
self.clss = None
self.track_history = defaultdict(list)
# Region & Line Information
self.counting_region = None
self.line_dist_thresh = line_dist_thresh
self.region_thickness = region_thickness
self.region_color = count_reg_color
# Object Counting Information
self.in_counts = 0
self.out_counts = 0
self.count_ids = []
self.class_wise_count = {}
self.count_txt_color = count_txt_color
self.count_bg_color = count_bg_color
self.cls_txtdisplay_gap = 50
# Decay factor
self.decay_factor = decay_factor
# Check if environment supports imshow
self.env_check = check_imshow(warn=True)
# Region and line selection
if count_reg_pts is not None:
if len(count_reg_pts) == 2:
self.count_reg_pts = count_reg_pts
print(self.count_reg_pts)
if self.count_reg_pts is not None:
if len(self.count_reg_pts) == 2:
print("Line Counter Initiated.")
self.count_reg_pts = count_reg_pts
self.counting_region = LineString(self.count_reg_pts)
elif len(count_reg_pts) >= 3:
elif len(self.count_reg_pts) >= 3:
print("Polygon Counter Initiated.")
self.count_reg_pts = count_reg_pts
self.counting_region = Polygon(self.count_reg_pts)
else:
print("Invalid Region points provided, region_points must be 2 for lines or >= 3 for polygons.")
print("Using Line Counter Now")
self.counting_region = LineString(self.count_reg_pts)
# Heatmap new frame
self.heatmap = np.zeros((int(self.imh), int(self.imw)), dtype=np.float32)
self.count_txt_color = count_txt_color
self.count_bg_color = count_bg_color
self.region_color = count_reg_color
self.region_thickness = region_thickness
self.decay_factor = decay_factor
self.line_dist_thresh = line_dist_thresh
self.shape = shape
# shape of heatmap, if not selected
# Shape of heatmap, if not selected
if self.shape not in {"circle", "rect"}:
print("Unknown shape value provided, 'circle' & 'rect' supported")
print("Using Circular shape now")
self.shape = "circle"
def extract_results(self, tracks):
def extract_results(self, tracks, _intialized=False):
"""
Extracts results from the provided data.
@ -171,18 +127,20 @@ class Heatmap:
tracks (list): List of tracks obtained from the object tracking process.
"""
self.im0 = im0
if tracks[0].boxes.id is None:
self.heatmap = np.zeros((int(self.imh), int(self.imw)), dtype=np.float32)
if self.view_img and self.env_check:
self.display_frames()
return im0
# Initialize heatmap only once
if not self.initialized:
self.heatmap = np.zeros((int(self.im0.shape[0]), int(self.im0.shape[1])), dtype=np.float32)
self.initialized = True
self.heatmap *= self.decay_factor # decay factor
self.extract_results(tracks)
self.annotator = Annotator(self.im0, self.tf, None)
if self.count_reg_pts is not None:
if self.track_ids is not None:
# Draw counting region
if self.view_in_counts or self.view_out_counts:
if self.count_reg_pts is not None:
self.annotator.draw_region(
reg_pts=self.count_reg_pts, color=self.region_color, thickness=self.region_thickness
)
@ -214,25 +172,12 @@ class Heatmap:
prev_position = self.track_history[track_id][-2] if len(self.track_history[track_id]) > 1 else None
# Count objects in any polygon
if len(self.count_reg_pts) >= 3:
is_inside = self.counting_region.contains(Point(track_line[-1]))
if self.count_reg_pts is not None:
# Count objects in any polygon
if len(self.count_reg_pts) >= 3:
is_inside = self.counting_region.contains(Point(track_line[-1]))
if prev_position is not None and is_inside and track_id not in self.count_ids:
self.count_ids.append(track_id)
if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
self.in_counts += 1
self.class_wise_count[self.names[cls]]["IN"] += 1
else:
self.out_counts += 1
self.class_wise_count[self.names[cls]]["OUT"] += 1
# Count objects using line
elif len(self.count_reg_pts) == 2:
if prev_position is not None and track_id not in self.count_ids:
distance = Point(track_line[-1]).distance(self.counting_region)
if distance < self.line_dist_thresh and track_id not in self.count_ids:
if prev_position is not None and is_inside and track_id not in self.count_ids:
self.count_ids.append(track_id)
if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
@ -242,6 +187,22 @@ class Heatmap:
self.out_counts += 1
self.class_wise_count[self.names[cls]]["OUT"] += 1
# Count objects using line
elif len(self.count_reg_pts) == 2:
if prev_position is not None and track_id not in self.count_ids:
distance = Point(track_line[-1]).distance(self.counting_region)
if distance < self.line_dist_thresh and track_id not in self.count_ids:
self.count_ids.append(track_id)
if (box[0] - prev_position[0]) * (
self.counting_region.centroid.x - prev_position[0]
) > 0:
self.in_counts += 1
self.class_wise_count[self.names[cls]]["IN"] += 1
else:
self.out_counts += 1
self.class_wise_count[self.names[cls]]["OUT"] += 1
else:
for box, cls in zip(self.boxes, self.clss):
if self.shape == "circle":
@ -258,26 +219,26 @@ class Heatmap:
else:
self.heatmap[int(box[1]) : int(box[3]), int(box[0]) : int(box[2])] += 2
if self.count_reg_pts is not None:
labels_dict = {}
for key, value in self.class_wise_count.items():
if value["IN"] != 0 or value["OUT"] != 0:
if not self.view_in_counts and not self.view_out_counts:
continue
elif not self.view_in_counts:
labels_dict[str.capitalize(key)] = f"OUT {value['OUT']}"
elif not self.view_out_counts:
labels_dict[str.capitalize(key)] = f"IN {value['IN']}"
else:
labels_dict[str.capitalize(key)] = f"IN {value['IN']} OUT {value['OUT']}"
if labels_dict is not None:
self.annotator.display_analytics(self.im0, labels_dict, self.count_txt_color, self.count_bg_color, 10)
# Normalize, apply colormap to heatmap and combine with original image
heatmap_normalized = cv2.normalize(self.heatmap, None, 0, 255, cv2.NORM_MINMAX)
heatmap_colored = cv2.applyColorMap(heatmap_normalized.astype(np.uint8), self.colormap)
labels_dict = {}
for key, value in self.class_wise_count.items():
if value["IN"] != 0 or value["OUT"] != 0:
if not self.view_in_counts and not self.view_out_counts:
continue
elif not self.view_in_counts:
labels_dict[str.capitalize(key)] = f"OUT {value['OUT']}"
elif not self.view_out_counts:
labels_dict[str.capitalize(key)] = f"IN {value['IN']}"
else:
labels_dict[str.capitalize(key)] = f"IN {value['IN']} OUT {value['OUT']}"
if labels_dict is not None:
self.annotator.display_analytics(self.im0, labels_dict, self.count_txt_color, self.count_bg_color, 10)
self.im0 = cv2.addWeighted(self.im0, 1 - self.heatmap_alpha, heatmap_colored, self.heatmap_alpha, 0)
if self.env_check and self.view_img:
@ -294,4 +255,5 @@ class Heatmap:
if __name__ == "__main__":
Heatmap()
classes_names = {0: "person", 1: "car"} # example class names
heatmap = Heatmap(classes_names)

144
ultralytics/solutions/object_counter.py
View file

@ -15,55 +15,10 @@ from shapely.geometry import LineString, Point, Polygon
class ObjectCounter:
"""A class to manage the counting of objects in a real-time video stream based on their tracks."""
def __init__(self):
"""Initializes the Counter with default values for various tracking and counting parameters."""
# Mouse events
self.is_drawing = False
self.selected_point = None
# Region & Line Information
self.reg_pts = [(20, 400), (1260, 400)]
self.line_dist_thresh = 15
self.counting_region = None
self.region_color = (255, 0, 255)
self.region_thickness = 5
# Image and annotation Information
self.im0 = None
self.tf = None
self.view_img = False
self.view_in_counts = True
self.view_out_counts = True
self.names = None # Classes names
self.annotator = None # Annotator
self.window_name = "Ultralytics YOLOv8 Object Counter"
# Object counting Information
self.in_counts = 0
self.out_counts = 0
self.count_ids = []
self.class_wise_count = {}
self.count_txt_thickness = 0
self.count_txt_color = (255, 255, 255)
self.count_bg_color = (255, 255, 255)
self.cls_txtdisplay_gap = 50
self.fontsize = 0.6
# Tracks info
self.track_history = defaultdict(list)
self.track_thickness = 2
self.draw_tracks = False
self.track_color = None
# Check if environment support imshow
self.env_check = check_imshow(warn=True)
def set_args(
def __init__(
self,
classes_names,
reg_pts,
reg_pts=None,
count_reg_color=(255, 0, 255),
count_txt_color=(0, 0, 0),
count_bg_color=(255, 255, 255),
@ -79,66 +34,90 @@ class ObjectCounter:
cls_txtdisplay_gap=50,
):
"""
Configures the Counter's image, bounding box line thickness, and counting region points.
Initializes the ObjectCounter with various tracking and counting parameters.
Args:
classes_names (dict): Dictionary of class names.
reg_pts (list): List of points defining the counting region.
count_reg_color (tuple): RGB color of the counting region.
count_txt_color (tuple): RGB color of the count text.
count_bg_color (tuple): RGB color of the count text background.
line_thickness (int): Line thickness for bounding boxes.
track_thickness (int): Thickness of the track lines.
view_img (bool): Flag to control whether to display the video stream.
view_in_counts (bool): Flag to control whether to display the incounts on video stream.
view_out_counts (bool): Flag to control whether to display the outcounts on video stream.
reg_pts (list): Initial list of points defining the counting region.
classes_names (dict): Classes names
track_thickness (int): Track thickness
draw_tracks (Bool): draw tracks
count_txt_color (RGB color): count text color value
count_bg_color (RGB color): count highlighter line color
count_reg_color (RGB color): Color of object counting region
track_color (RGB color): color for tracks
region_thickness (int): Object counting Region thickness
line_dist_thresh (int): Euclidean Distance threshold for line counter
cls_txtdisplay_gap (int): Display gap between each class count
view_in_counts (bool): Flag to control whether to display the in counts on the video stream.
view_out_counts (bool): Flag to control whether to display the out counts on the video stream.
draw_tracks (bool): Flag to control whether to draw the object tracks.
track_color (tuple): RGB color of the tracks.
region_thickness (int): Thickness of the object counting region.
line_dist_thresh (int): Euclidean distance threshold for line counter.
cls_txtdisplay_gap (int): Display gap between each class count.
"""
# Mouse events
self.is_drawing = False
self.selected_point = None
# Region & Line Information
self.reg_pts = [(20, 400), (1260, 400)] if reg_pts is None else reg_pts
self.line_dist_thresh = line_dist_thresh
self.counting_region = None
self.region_color = count_reg_color
self.region_thickness = region_thickness
# Image and annotation Information
self.im0 = None
self.tf = line_thickness
self.view_img = view_img
self.view_in_counts = view_in_counts
self.view_out_counts = view_out_counts
self.names = classes_names # Classes names
self.annotator = None # Annotator
self.window_name = "Ultralytics YOLOv8 Object Counter"
# Object counting Information
self.in_counts = 0
self.out_counts = 0
self.count_ids = []
self.class_wise_count = {}
self.count_txt_thickness = 0
self.count_txt_color = count_txt_color
self.count_bg_color = count_bg_color
self.cls_txtdisplay_gap = cls_txtdisplay_gap
self.fontsize = 0.6
# Tracks info
self.track_history = defaultdict(list)
self.track_thickness = track_thickness
self.draw_tracks = draw_tracks
self.track_color = track_color
# Region and line selection
if len(reg_pts) == 2:
# Check if environment supports imshow
self.env_check = check_imshow(warn=True)
# Initialize counting region
if len(self.reg_pts) == 2:
print("Line Counter Initiated.")
self.reg_pts = reg_pts
self.counting_region = LineString(self.reg_pts)
elif len(reg_pts) >= 3:
elif len(self.reg_pts) >= 3:
print("Polygon Counter Initiated.")
self.reg_pts = reg_pts
self.counting_region = Polygon(self.reg_pts)
else:
print("Invalid Region points provided, region_points must be 2 for lines or >= 3 for polygons.")
print("Using Line Counter Now")
self.counting_region = LineString(self.reg_pts)
self.names = classes_names
self.track_color = track_color
self.count_txt_color = count_txt_color
self.count_bg_color = count_bg_color
self.region_color = count_reg_color
self.region_thickness = region_thickness
self.line_dist_thresh = line_dist_thresh
self.cls_txtdisplay_gap = cls_txtdisplay_gap
def mouse_event_for_region(self, event, x, y, flags, params):
"""
This function is designed to move region with mouse events in a real-time video stream.
Handles mouse events for defining and moving the counting region 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.).
params (dict): Additional parameters you may want to pass to the function.
flags (int): Any associated event flags (e.g., cv2.EVENT_FLAG_CTRLKEY, cv2.EVENT_FLAG_SHIFTKEY, etc.).
params (dict): Additional parameters for the function.
"""
if event == cv2.EVENT_LBUTTONDOWN:
for i, point in enumerate(self.reg_pts):
@ -240,11 +219,11 @@ class ObjectCounter:
else:
labels_dict[str.capitalize(key)] = f"IN {value['IN']} OUT {value['OUT']}"
if labels_dict is not None:
if labels_dict:
self.annotator.display_analytics(self.im0, labels_dict, self.count_txt_color, self.count_bg_color, 10)
def display_frames(self):
"""Display frame."""
"""Displays the current frame with annotations and regions in a window."""
if self.env_check:
cv2.namedWindow(self.window_name)
if len(self.reg_pts) == 4: # only add mouse event If user drawn region
@ -271,4 +250,5 @@ class ObjectCounter:
if __name__ == "__main__":
ObjectCounter()
classes_names = {0: "person", 1: "car"} # example class names
ObjectCounter(classes_names)

31
ultralytics/solutions/parking_management.py
View file

@ -8,17 +8,22 @@ from PIL import Image, ImageTk
from ultralytics.utils.checks import check_imshow, check_requirements
from ultralytics.utils.plotting import Annotator
check_requirements("tkinter")
import tkinter as tk
class ParkingPtsSelection:
def __init__(self, master):
"""Initializes the UI for selecting parking zone points in a tkinter window."""
"""
Initializes the UI for selecting parking zone points in a tkinter window.
Args:
master (tk.Tk): The main tkinter window object.
"""
check_requirements("tkinter")
import tkinter as tk
self.master = master
master.title("Ultralytics Parking Zones Points Selector")
# Resizable false
# Disable window resizing
master.resizable(False, False)
# Setup canvas for image display
@ -36,7 +41,6 @@ class ParkingPtsSelection:
self.image_path = None
self.image = None
self.canvas_image = None
self.canvas = None
self.bounding_boxes = []
self.current_box = []
self.img_width = 0
@ -101,7 +105,6 @@ class ParkingPtsSelection:
Args:
box (list): Bounding box data
"""
for i in range(4):
x1, y1 = box[i]
x2, y2 = box[(i + 1) % 4]
@ -151,6 +154,17 @@ class ParkingManagement:
available_region_color=(0, 0, 255),
margin=10,
):
"""
Initializes the parking management system with a YOLOv8 model and visualization settings.
Args:
model_path (str): Path to the YOLOv8 model.
txt_color (tuple): RGB color tuple for text.
bg_color (tuple): RGB color tuple for background.
occupied_region_color (tuple): RGB color tuple for occupied regions.
available_region_color (tuple): RGB color tuple for available regions.
margin (int): Margin for text display.
"""
# Model path and initialization
self.model_path = model_path
self.model = self.load_model()
@ -166,7 +180,7 @@ class ParkingManagement:
self.available_region_color = available_region_color
self.window_name = "Ultralytics YOLOv8 Parking Management System"
# Check if environment support imshow
# Check if environment supports imshow
self.env_check = check_imshow(warn=True)
def load_model(self):
@ -184,7 +198,6 @@ class ParkingManagement:
Args:
json_file (str): file that have all parking slot points
"""
with open(json_file, "r") as json_file:
return json.load(json_file)

145
ultralytics/solutions/queue_management.py
View file

@ -13,49 +13,12 @@ from shapely.geometry import Point, Polygon
class QueueManager:
"""A class to manage the queue management in real-time video stream based on their tracks."""
"""A class to manage the queue in a real-time video stream based on object tracks."""
def __init__(self):
"""Initializes the queue manager with default values for various tracking and counting parameters."""
# Mouse events
self.is_drawing = False
self.selected_point = None
# Region & Line Information
self.reg_pts = [(20, 60), (20, 680), (1120, 680), (1120, 60)]
self.counting_region = None
self.region_color = (255, 0, 255)
self.region_thickness = 5
# Image and annotation Information
self.im0 = None
self.tf = None
self.view_img = False
self.view_queue_counts = True
self.fontsize = 0.6
self.names = None # Classes names
self.annotator = None # Annotator
self.window_name = "Ultralytics YOLOv8 Queue Manager"
# Object counting Information
self.counts = 0
self.count_txt_color = (255, 255, 255)
# Tracks info
self.track_history = defaultdict(list)
self.track_thickness = 2
self.draw_tracks = False
self.track_color = None
# Check if environment support imshow
self.env_check = check_imshow(warn=True)
def set_args(
def __init__(
self,
classes_names,
reg_pts,
reg_pts=None,
line_thickness=2,
track_thickness=2,
view_img=False,
@ -68,48 +31,65 @@ class QueueManager:
fontsize=0.7,
):
"""
Configures the Counter's image, bounding box line thickness, and counting region points.
Initializes the QueueManager with specified parameters for tracking and counting objects.
Args:
line_thickness (int): Line thickness for bounding boxes.
view_img (bool): Flag to control whether to display the video stream.
view_queue_counts (bool): Flag to control whether to display the counts on video stream.
reg_pts (list): Initial list of points defining the counting region.
classes_names (dict): Classes names
region_color (RGB color): Color of queue region
track_thickness (int): Track thickness
draw_tracks (Bool): draw tracks
count_txt_color (RGB color): count text color value
track_color (RGB color): color for tracks
region_thickness (int): Object counting Region thickness
fontsize (float): Text display font size
classes_names (dict): A dictionary mapping class IDs to class names.
reg_pts (list of tuples, optional): Points defining the counting region polygon. Defaults to a predefined
rectangle.
line_thickness (int, optional): Thickness of the annotation lines. Defaults to 2.
track_thickness (int, optional): Thickness of the track lines. Defaults to 2.
view_img (bool, optional): Whether to display the image frames. Defaults to False.
region_color (tuple, optional): Color of the counting region lines (BGR). Defaults to (255, 0, 255).
view_queue_counts (bool, optional): Whether to display the queue counts. Defaults to True.
draw_tracks (bool, optional): Whether to draw tracks of the objects. Defaults to False.
count_txt_color (tuple, optional): Color of the count text (BGR). Defaults to (255, 255, 255).
track_color (tuple, optional): Color of the tracks. If None, different colors will be used for different
tracks. Defaults to None.
region_thickness (int, optional): Thickness of the counting region lines. Defaults to 5.
fontsize (float, optional): Font size for the text annotations. Defaults to 0.7.
"""
# Mouse events state
self.is_drawing = False
self.selected_point = None
# Region & Line Information
self.reg_pts = reg_pts if reg_pts is not None else [(20, 60), (20, 680), (1120, 680), (1120, 60)]
self.counting_region = (
Polygon(self.reg_pts) if len(self.reg_pts) >= 3 else Polygon([(20, 60), (20, 680), (1120, 680), (1120, 60)])
)
self.region_color = region_color
self.region_thickness = region_thickness
# Image and annotation Information
self.im0 = None
self.tf = line_thickness
self.view_img = view_img
self.view_queue_counts = view_queue_counts
self.fontsize = fontsize
self.names = classes_names # Class names
self.annotator = None # Annotator
self.window_name = "Ultralytics YOLOv8 Queue Manager"
# Object counting Information
self.counts = 0
self.count_txt_color = count_txt_color
# Tracks info
self.track_history = defaultdict(list)
self.track_thickness = track_thickness
self.draw_tracks = draw_tracks
self.region_color = region_color
if len(reg_pts) >= 3:
print("Queue region initiated...")
self.reg_pts = reg_pts
self.counting_region = Polygon(self.reg_pts)
else:
print("Invalid region points provided...")
print("Using default region now....")
self.counting_region = Polygon(self.reg_pts)
self.names = classes_names
self.track_color = track_color
self.count_txt_color = count_txt_color
self.region_thickness = region_thickness
self.fontsize = fontsize
# Check if environment supports imshow
self.env_check = check_imshow(warn=True)
def extract_and_process_tracks(self, tracks):
"""Extracts and processes tracks for queue management in a video stream."""
# Annotator Init and queue region drawing
# Initialize annotator and draw the queue region
self.annotator = Annotator(self.im0, self.tf, self.names)
if tracks[0].boxes.id is not None:
@ -122,48 +102,48 @@ class QueueManager:
# Draw bounding box
self.annotator.box_label(box, label=f"{self.names[cls]}#{track_id}", color=colors(int(track_id), True))
# Draw Tracks
# Update track history
track_line = self.track_history[track_id]
track_line.append((float((box[0] + box[2]) / 2), float((box[1] + box[3]) / 2)))
if len(track_line) > 30:
track_line.pop(0)
# Draw track trails
# Draw track trails if enabled
if self.draw_tracks:
self.annotator.draw_centroid_and_tracks(
track_line,
color=self.track_color if self.track_color else colors(int(track_id), True),
color=self.track_color or colors(int(track_id), True),
track_thickness=self.track_thickness,
)
prev_position = self.track_history[track_id][-2] if len(self.track_history[track_id]) > 1 else None
# Check if the object is inside the counting region
if len(self.reg_pts) >= 3:
is_inside = self.counting_region.contains(Point(track_line[-1]))
if prev_position is not None and is_inside:
self.counts += 1
label = "Queue Counts : " + str(self.counts)
# Display queue counts
label = f"Queue Counts : {str(self.counts)}"
if label is not None:
self.annotator.queue_counts_display(
label,
points=self.reg_pts,
region_color=self.region_color,
txt_color=self.count_txt_color,
fontsize=self.fontsize,
)
self.counts = 0
self.counts = 0 # Reset counts after displaying
self.display_frames()
def display_frames(self):
"""Display frame."""
"""Displays the current frame with annotations."""
if self.env_check:
self.annotator.draw_region(reg_pts=self.reg_pts, thickness=self.region_thickness, color=self.region_color)
cv2.namedWindow(self.window_name)
cv2.imshow(self.window_name, self.im0)
# Break Window
# Close window on 'q' key press
if cv2.waitKey(1) & 0xFF == ord("q"):
return
@ -175,13 +155,14 @@ class QueueManager:
im0 (ndarray): Current frame from the video stream.
tracks (list): List of tracks obtained from the object tracking process.
"""
self.im0 = im0 # store image
self.extract_and_process_tracks(tracks) # draw region even if no objects
self.im0 = im0 # Store the current frame
self.extract_and_process_tracks(tracks) # Extract and process tracks
if self.view_img:
self.display_frames()
self.display_frames() # Display the frame if enabled
return self.im0
if __name__ == "__main__":
QueueManager()
classes_names = {0: "person", 1: "car"} # example class names
queue_manager = QueueManager(classes_names)

118
ultralytics/solutions/speed_estimation.py
View file

@ -11,73 +11,52 @@ from ultralytics.utils.plotting import Annotator, colors
class SpeedEstimator:
"""A class to estimation speed of objects in real-time video stream based on their tracks."""
"""A class to estimate the speed of objects in a real-time video stream based on their tracks."""
def __init__(self):
"""Initializes the speed-estimator class with default values for Visual, Image, track and speed parameters."""
def __init__(self, names, reg_pts=None, view_img=False, line_thickness=2, region_thickness=5, spdl_dist_thresh=10):
"""
Initializes the SpeedEstimator with the given parameters.
# Visual & im0 information
Args:
names (dict): Dictionary of class names.
reg_pts (list, optional): List of region points for speed estimation. Defaults to [(20, 400), (1260, 400)].
view_img (bool, optional): Whether to display the image with annotations. Defaults to False.
line_thickness (int, optional): Thickness of the lines for drawing boxes and tracks. Defaults to 2.
region_thickness (int, optional): Thickness of the region lines. Defaults to 5.
spdl_dist_thresh (int, optional): Distance threshold for speed calculation. Defaults to 10.
"""
# Visual & image information
self.im0 = None
self.annotator = None
self.view_img = False
self.view_img = view_img
# Region information
self.reg_pts = [(20, 400), (1260, 400)]
self.region_thickness = 3
self.reg_pts = reg_pts if reg_pts is not None else [(20, 400), (1260, 400)]
self.region_thickness = region_thickness
# Predict/track information
# Tracking information
self.clss = None
self.names = None
self.names = names
self.boxes = None
self.trk_ids = None
self.trk_pts = None
self.line_thickness = 2
self.line_thickness = line_thickness
self.trk_history = defaultdict(list)
# Speed estimator information
# Speed estimation information
self.current_time = 0
self.dist_data = {}
self.trk_idslist = []
self.spdl_dist_thresh = 10
self.spdl_dist_thresh = spdl_dist_thresh
self.trk_previous_times = {}
self.trk_previous_points = {}
# Check if environment support imshow
# Check if the environment supports imshow
self.env_check = check_imshow(warn=True)
def set_args(
self,
reg_pts,
names,
view_img=False,
line_thickness=2,
region_thickness=5,
spdl_dist_thresh=10,
):
"""
Configures the speed estimation and display parameters.
Args:
reg_pts (list): Initial list of points defining the speed calculation region.
names (dict): object detection classes names
view_img (bool): Flag indicating frame display
line_thickness (int): Line thickness for bounding boxes.
region_thickness (int): Speed estimation region thickness
spdl_dist_thresh (int): Euclidean distance threshold for speed line
"""
if reg_pts is None:
print("Region points not provided, using default values")
else:
self.reg_pts = reg_pts
self.names = names
self.view_img = view_img
self.line_thickness = line_thickness
self.region_thickness = region_thickness
self.spdl_dist_thresh = spdl_dist_thresh
def extract_tracks(self, tracks):
"""
Extracts results from the provided data.
Extracts results from the provided tracking data.
Args:
tracks (list): List of tracks obtained from the object tracking process.
@ -88,11 +67,14 @@ class SpeedEstimator:
def store_track_info(self, track_id, box):
"""
Store track data.
Stores track data.
Args:
track_id (int): object track id.
box (list): object bounding box data
track_id (int): Object track id.
box (list): Object bounding box data.
Returns:
(list): Updated tracking history for the given track_id.
"""
track = self.trk_history[track_id]
bbox_center = (float((box[0] + box[2]) / 2), float((box[1] + box[3]) / 2))
@ -106,43 +88,39 @@ class SpeedEstimator:
def plot_box_and_track(self, track_id, box, cls, track):
"""
Plot track and bounding box.
Plots track and bounding box.
Args:
track_id (int): object track id.
box (list): object bounding box data
cls (str): object class name
track (list): tracking history for tracks path drawing
track_id (int): Object track id.
box (list): Object bounding box data.
cls (str): Object class name.
track (list): Tracking history for drawing tracks path.
"""
speed_label = f"{int(self.dist_data[track_id])}km/ph" if track_id in self.dist_data else self.names[int(cls)]
speed_label = f"{int(self.dist_data[track_id])} km/h" if track_id in self.dist_data else self.names[int(cls)]
bbox_color = colors(int(track_id)) if track_id in self.dist_data else (255, 0, 255)
self.annotator.box_label(box, speed_label, bbox_color)
cv2.polylines(self.im0, [self.trk_pts], isClosed=False, color=(0, 255, 0), thickness=1)
cv2.circle(self.im0, (int(track[-1][0]), int(track[-1][1])), 5, bbox_color, -1)
def calculate_speed(self, trk_id, track):
"""
Calculation of object speed.
Calculates the speed of an object.
Args:
trk_id (int): object track id.
track (list): tracking history for tracks path drawing
trk_id (int): Object track id.
track (list): Tracking history for drawing tracks path.
"""
if not self.reg_pts[0][0] < track[-1][0] < self.reg_pts[1][0]:
return
if self.reg_pts[1][1] - self.spdl_dist_thresh < track[-1][1] < self.reg_pts[1][1] + self.spdl_dist_thresh:
direction = "known"
elif self.reg_pts[0][1] - self.spdl_dist_thresh < track[-1][1] < self.reg_pts[0][1] + self.spdl_dist_thresh:
direction = "known"
else:
direction = "unknown"
if self.trk_previous_times[trk_id] != 0 and direction != "unknown" and trk_id not in self.trk_idslist:
if self.trk_previous_times.get(trk_id) != 0 and direction != "unknown" and trk_id not in self.trk_idslist:
self.trk_idslist.append(trk_id)
time_difference = time() - self.trk_previous_times[trk_id]
@ -156,21 +134,24 @@ class SpeedEstimator:
def estimate_speed(self, im0, tracks, region_color=(255, 0, 0)):
"""
Calculate object based on tracking data.
Estimates the speed of objects based on tracking data.
Args:
im0 (nd array): Image
im0 (ndarray): Image.
tracks (list): List of tracks obtained from the object tracking process.
region_color (tuple): Color to use when drawing regions.
region_color (tuple, optional): Color to use when drawing regions. Defaults to (255, 0, 0).
Returns:
(ndarray): The image with annotated boxes and tracks.
"""
self.im0 = im0
if tracks[0].boxes.id is None:
if self.view_img and self.env_check:
self.display_frames()
return im0
self.extract_tracks(tracks)
self.annotator = Annotator(self.im0, line_width=2)
self.extract_tracks(tracks)
self.annotator = Annotator(self.im0, line_width=self.line_thickness)
self.annotator.draw_region(reg_pts=self.reg_pts, color=region_color, thickness=self.region_thickness)
for box, trk_id, cls in zip(self.boxes, self.trk_ids, self.clss):
@ -188,11 +169,12 @@ class SpeedEstimator:
return im0
def display_frames(self):
"""Display frame."""
"""Displays the current frame."""
cv2.imshow("Ultralytics Speed Estimation", self.im0)
if cv2.waitKey(1) & 0xFF == ord("q"):
return
if __name__ == "__main__":
SpeedEstimator()
names = {0: "person", 1: "car"} # example class names
speed_estimator = SpeedEstimator(names)