Yolov8 源码解析（三十八）
.\yolov8\ultralytics\nn\__init__.py
# 导入模块中的特定类和函数，包括：
# - BaseModel: 基础模型类
# - ClassificationModel: 分类模型类
# - DetectionModel: 目标检测模型类
# - SegmentationModel: 分割模型类
# - attempt_load_one_weight: 尝试加载单个权重函数
# - attempt_load_weights: 尝试加载权重函数
# - guess_model_scale: 推测模型规模函数
# - guess_model_task: 推测模型任务函数
# - parse_model: 解析模型函数
# - torch_safe_load: 安全加载 Torch 模型函数
# - yaml_model_load: 加载 YAML 格式模型函数
from .tasks import (
BaseModel,
ClassificationModel,
DetectionModel,
SegmentationModel,
attempt_load_one_weight,
attempt_load_weights,
guess_model_scale,
guess_model_task,
parse_model,
torch_safe_load,
yaml_model_load,
)
# 模块中可以直接访问的全部对象的元组，包括类和函数
__all__ = (
"attempt_load_one_weight",
"attempt_load_weights",
"parse_model",
"yaml_model_load",
"guess_model_task",
"guess_model_scale",
"torch_safe_load",
"DetectionModel",
"SegmentationModel",
"ClassificationModel",
"BaseModel",
)
.\yolov8\ultralytics\solutions\ai_gym.py
# 导入OpenCV库，用于图像处理
import cv2
# 导入自定义函数和类
from ultralytics.utils.checks import check_imshow
from ultralytics.utils.plotting import Annotator
# AIGym类用于实时视频流中人员姿势的管理
class AIGym:
"""A class to manage the gym steps of people in a real-time video stream based on their poses."""
def __init__(
self,
kpts_to_check,
line_thickness=2,
view_img=False,
pose_up_angle=145.0,
pose_down_angle=90.0,
pose_type="pullup",
):
"""
Initializes the AIGym class with the specified parameters.
Args:
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.im0 = None # 初始图像设为None
self.tf = line_thickness # 线条厚度设定为传入的参数值
# 关键点和计数信息
self.keypoints = None # 关键点初始化为None
self.poseup_angle = pose_up_angle # 'up'姿势的角度阈值
self.posedown_angle = pose_down_angle # 'down'姿势的角度阈值
self.threshold = 0.001 # 阈值设定为0.001
# 存储阶段、计数和角度信息
self.angle = None # 角度信息初始化为None
self.count = None # 计数信息初始化为None
self.stage = None # 阶段信息初始化为None
self.pose_type = pose_type # 姿势类型，默认为"pullup"
self.kpts_to_check = kpts_to_check # 需要检查的关键点索引列表
# 可视化信息
self.view_img = view_img # 是否显示图像的标志
self.annotator = None # 标注器初始化为None
# 检查环境是否支持imshow函数
self.env_check = check_imshow(warn=True) # 调用自定义函数检查环境支持情况
self.count = [] # 计数列表初始化为空列表
self.angle = [] # 角度列表初始化为空列表
self.stage = [] # 阶段列表初始化为空列表
def start_counting(self, im0, results):
"""
Function used to count the gym steps.
Args:
im0 (ndarray): Current frame from the video stream.
results (list): Pose estimation data.
"""
# 将当前帧图像保存到对象的属性中
self.im0 = im0
# 如果没有检测到姿态估计数据，则直接返回原始图像
if not len(results[0]):
return self.im0
# 如果检测到的人数超过已记录的计数器数量，进行扩展
if len(results[0]) > len(self.count):
new_human = len(results[0]) - len(self.count)
self.count += [0] * new_human
self.angle += [0] * new_human
self.stage += ["-"] * new_human
# 获取关键点数据
self.keypoints = results[0].keypoints.data
# 创建一个用于绘制的注释器对象
self.annotator = Annotator(im0, line_width=self.tf)
# 遍历检测到的关键点数据
for ind, k in enumerate(reversed(self.keypoints)):
# 估算姿势角度并绘制特定关键点
if self.pose_type in {"pushup", "pullup", "abworkout", "squat"}:
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.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 in {"pushup", "squat"}:
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])],
)
# 绘制关键点
self.annotator.kpts(k, shape=(640, 640), radius=1, kpt_line=True)
# 如果环境支持并且需要显示图像，则显示处理后的图像
if self.env_check and self.view_img:
cv2.imshow("Ultralytics YOLOv8 AI GYM", self.im0)
# 等待用户按键以退出显示
if cv2.waitKey(1) & 0xFF == ord("q"):
return
# 返回处理后的图像
return self.im0
# 如果当前脚本作为主程序运行（而非被导入其他模块），执行以下代码块
if __name__ == "__main__":
# 定义一个示例的关键点列表，用于检查
kpts_to_check = [0, 1, 2] # example keypoints
# 创建一个 AIGym 对象，并传入关键点列表作为参数
aigym = AIGym(kpts_to_check)
.\yolov8\ultralytics\solutions\analytics.py
# 导入警告模块，用于处理警告信息
import warnings
# 导入循环迭代工具模块，用于创建迭代器
from itertools import cycle
# 导入OpenCV库，用于图像处理
import cv2
# 导入matplotlib.pyplot模块，用于绘制图表
import matplotlib.pyplot as plt
# 导入NumPy库，用于数值计算和数组操作
import numpy as np
# 导入matplotlib的FigureCanvas类，用于绘制图形
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
# 导入matplotlib的Figure类，用于创建图形对象
from matplotlib.figure import Figure
class Analytics:
"""一个用于创建和更新各种类型图表（线形图、柱状图、饼图、面积图）的类，用于视觉分析。"""
def __init__(
self,
type,
writer,
im0_shape,
title="ultralytics",
x_label="x",
y_label="y",
bg_color="white",
fg_color="black",
line_color="yellow",
line_width=2,
points_width=10,
fontsize=13,
view_img=False,
save_img=True,
max_points=50,
def update_area(self, frame_number, counts_dict):
"""
Update the area graph with new data for multiple classes.
Args:
frame_number (int): The current frame number.
counts_dict (dict): Dictionary with class names as keys and counts as values.
"""
# 初始化 x_data 为空数组
x_data = np.array([])
# 初始化 y_data_dict，使用 counts_dict 的键创建对应的空数组作为值
y_data_dict = {key: np.array([]) for key in counts_dict.keys()}
# 如果图形已经存在线条
if self.ax.lines:
# 获取第一条线的 x 轴数据
x_data = self.ax.lines[0].get_xdata()
# 遍历每条线并更新对应类别的 y 轴数据
for line, key in zip(self.ax.lines, counts_dict.keys()):
y_data_dict[key] = line.get_ydata()
# 将当前帧数添加到 x_data 中
x_data = np.append(x_data, float(frame_number))
max_length = len(x_data)
# 遍历每个类别的数据
for key in counts_dict.keys():
# 将新的计数值添加到对应类别的 y 数据中
y_data_dict[key] = np.append(y_data_dict[key], float(counts_dict[key]))
# 如果某个类别的 y 数据长度小于 max_length，则用常数填充
if len(y_data_dict[key]) < max_length:
y_data_dict[key] = np.pad(y_data_dict[key], (0, max_length - len(y_data_dict[key])), "constant")
# 如果 x_data 的长度超过了 max_points，则移除最旧的点
if len(x_data) > self.max_points:
x_data = x_data[1:]
for key in counts_dict.keys():
y_data_dict[key] = y_data_dict[key][1:]
# 清空当前图形
self.ax.clear()
# 设置颜色循环使用的颜色列表
colors = ["#E1FF25", "#0BDBEB", "#FF64DA", "#111F68", "#042AFF"]
# 创建一个颜色循环迭代器
color_cycle = cycle(colors)
# 遍历每个类别及其对应的 y 数据
for key, y_data in y_data_dict.items():
# 获取下一个颜色
color = next(color_cycle)
# 填充区域图形
self.ax.fill_between(x_data, y_data, color=color, alpha=0.6)
# 绘制线条并设置线条属性
self.ax.plot(
x_data,
y_data,
color=color,
linewidth=self.line_width,
marker="o",
markersize=self.points_width,
label=f"{key} Data Points",
)
# 设置图形标题、x 轴标签和 y 轴标签的属性
self.ax.set_title(self.title, color=self.fg_color, fontsize=self.fontsize)
self.ax.set_xlabel(self.x_label, color=self.fg_color, fontsize=self.fontsize - 3)
self.ax.set_ylabel(self.y_label, color=self.fg_color, fontsize=self.fontsize - 3)
# 设置图例的位置、字体大小、背景颜色和边框颜色
legend = self.ax.legend(loc="upper left", fontsize=13, facecolor=self.bg_color, edgecolor=self.fg_color)
# 设置图例文本的颜色为前景色
for text in legend.get_texts():
text.set_color(self.fg_color)
# 绘制更新后的图形
self.canvas.draw()
# 将画布转换为 RGBA 缓冲区数组
im0 = np.array(self.canvas.renderer.buffer_rgba())
# 将图像数据写入并显示
self.write_and_display(im0)
def update_line(self, frame_number, total_counts):
"""
Update the line graph with new data.
Args:
frame_number (int): The current frame number.
total_counts (int): The total counts to plot.
"""
# 获取当前线图的 x 和 y 数据
x_data = self.line.get_xdata()
y_data = self.line.get_ydata()
# 将新的 frame_number 和 total_counts 添加到 x_data 和 y_data 中
x_data = np.append(x_data, float(frame_number))
y_data = np.append(y_data, float(total_counts))
# 更新线图的数据
self.line.set_data(x_data, y_data)
# 重新计算坐标轴限制
self.ax.relim()
# 自动调整视图范围
self.ax.autoscale_view()
# 重新绘制画布
self.canvas.draw()
# 将画布转换为 RGBA 缓冲区图像
im0 = np.array(self.canvas.renderer.buffer_rgba())
# 将图像写入并显示
self.write_and_display(im0)
def update_multiple_lines(self, counts_dict, labels_list, frame_number):
"""
Update the line graph with multiple classes.
Args:
counts_dict (int): Dictionary include each class counts.
labels_list (int): list include each classes names.
frame_number (int): The current frame number.
"""
# 发出警告，多条线的显示不受支持，将正常存储输出！
warnings.warn("Display is not supported for multiple lines, output will be stored normally!")
# 遍历所有标签
for obj in labels_list:
# 如果标签不在已存在的线图对象中，则创建新的线图对象
if obj not in self.lines:
(line,) = self.ax.plot([], [], label=obj, marker="o", markersize=self.points_width)
self.lines[obj] = line
# 获取当前标签对应的线图对象的 x 和 y 数据
x_data = self.lines[obj].get_xdata()
y_data = self.lines[obj].get_ydata()
# 如果数据点超过最大点数限制，则删除最早的数据点
if len(x_data) >= self.max_points:
x_data = np.delete(x_data, 0)
y_data = np.delete(y_data, 0)
# 将新的 frame_number 和对应类别的 counts 添加到 x_data 和 y_data 中
x_data = np.append(x_data, float(frame_number))
y_data = np.append(y_data, float(counts_dict.get(obj, 0)))
# 更新当前标签对应的线图对象的数据
self.lines[obj].set_data(x_data, y_data)
# 重新计算坐标轴限制
self.ax.relim()
# 自动调整视图范围
self.ax.autoscale_view()
# 添加图例
self.ax.legend()
# 重新绘制画布
self.canvas.draw()
# 将画布转换为 RGBA 缓冲区图像
im0 = np.array(self.canvas.renderer.buffer_rgba())
# 多条线的视图暂不支持，将 view_img 设置为 False
self.view_img = False # for multiple line view_img not supported yet, coming soon!
# 将图像写入并显示
self.write_and_display(im0)
def write_and_display(self, im0):
"""
Write and display the line graph
Args:
im0 (ndarray): Image for processing
"""
# 转换图像格式从 RGBA 到 BGR
im0 = cv2.cvtColor(im0[:, :, :3], cv2.COLOR_RGBA2BGR)
# 如果 view_img 为 True，则显示图像
cv2.imshow(self.title, im0) if self.view_img else None
# 如果 save_img 为 True，则写入图像
self.writer.write(im0) if self.save_img else None
def update_bar(self, count_dict):
"""
Update the bar graph with new data.
Args:
count_dict (dict): Dictionary containing the count data to plot.
"""
# 清空当前图形并设置背景颜色
self.ax.clear()
self.ax.set_facecolor(self.bg_color)
# 获取标签和计数数据
labels = list(count_dict.keys())
counts = list(count_dict.values())
# 将标签映射到颜色
for label in labels:
if label not in self.color_mapping:
self.color_mapping[label] = next(self.color_cycle)
colors = [self.color_mapping[label] for label in labels]
# 使用颜色绘制柱状图
bars = self.ax.bar(labels, counts, color=colors)
# 在柱状图上方显示数值
for bar, count in zip(bars, counts):
self.ax.text(
bar.get_x() + bar.get_width() / 2,
bar.get_height(),
str(count),
ha="center",
va="bottom",
color=self.fg_color,
)
# 显示和保存更新后的图形
canvas = FigureCanvas(self.fig)
canvas.draw()
buf = canvas.buffer_rgba()
im0 = np.asarray(buf)
self.write_and_display(im0)
def update_pie(self, classes_dict):
"""
Update the pie chart with new data.
Args:
classes_dict (dict): Dictionary containing the class data to plot.
"""
# 更新饼图数据
labels = list(classes_dict.keys())
sizes = list(classes_dict.values())
total = sum(sizes)
percentages = [size / total * 100 for size in sizes]
start_angle = 90
# 清空当前图形
self.ax.clear()
# 创建饼图，并设置起始角度及文本颜色
wedges, autotexts = self.ax.pie(sizes, autopct=None, startangle=start_angle, textprops={"color": self.fg_color})
# 构建带百分比的图例标签
legend_labels = [f"{label} ({percentage:.1f}%)" for label, percentage in zip(labels, percentages)]
self.ax.legend(wedges, legend_labels, title="Classes", loc="center left", bbox_to_anchor=(1, 0, 0.5, 1))
# 调整布局以适应图例
self.fig.tight_layout()
self.fig.subplots_adjust(left=0.1, right=0.75)
# 显示和保存更新后的饼图
im0 = self.fig.canvas.draw()
im0 = np.array(self.fig.canvas.renderer.buffer_rgba())
self.write_and_display(im0)
# 如果脚本被直接执行（而不是被导入为模块），则执行以下代码块
if __name__ == "__main__":
# 创建一个 Analytics 对象，设置参数为 "line"，writer 为 None，im0_shape 为 None
Analytics("line", writer=None, im0_shape=None)
.\yolov8\ultralytics\solutions\distance_calculation.py
# Ultralytics YOLO 🚀, AGPL-3.0 license
# 导入数学库
import math
# 导入 OpenCV 库
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 a real-time video stream based on their tracks."""
def __init__(
self,
names,
pixels_per_meter=10,
view_img=False,
line_thickness=2,
line_color=(255, 255, 0),
centroid_color=(255, 0, 255),
):
"""
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).
centroid_color (tuple, optional): Color of the centroids drawn (BGR format). Defaults to (255, 0, 255).
"""
# 图像和注解器相关信息初始化
self.im0 = None # 初始图像置空
self.annotator = None # 注解器置空
self.view_img = view_img # 是否显示视频流
self.line_color = line_color # 线条颜色
self.centroid_color = centroid_color # 质心颜色
# 预测和跟踪信息初始化
self.clss = None # 类别信息置空
self.names = names # 类别名称字典
self.boxes = None # 边界框信息置空
self.line_thickness = line_thickness # 线条粗细
self.trk_ids = None # 跟踪 ID 信息置空
# 距离计算信息初始化
self.centroids = [] # 质心列表
self.pixel_per_meter = pixels_per_meter # 像素与米的转换因子
# 鼠标事件信息初始化
self.left_mouse_count = 0 # 左键点击次数
self.selected_boxes = {} # 选中的边界框字典
# 检查环境是否支持 imshow 函数
self.env_check = check_imshow(warn=True)
# 处理鼠标事件以选择实时视频流中的区域
def mouse_event_for_distance(self, event, x, y, flags, param):
"""
Handles mouse events to select regions in a real-time video stream.
Args:
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.
"""
# 如果是左键单击事件
if event == cv2.EVENT_LBUTTONDOWN:
# 增加左键点击计数
self.left_mouse_count += 1
# 如果左键点击次数小于等于2
if self.left_mouse_count <= 2:
# 遍历每个盒子和其对应的跟踪 ID
for box, track_id in zip(self.boxes, self.trk_ids):
# 如果鼠标点击在当前盒子的范围内，并且该跟踪 ID 不在已选择的盒子中
if box[0] < x < box[2] and box[1] < y < box[3] and track_id not in self.selected_boxes:
# 将该跟踪 ID 和盒子加入已选择的盒子字典中
self.selected_boxes[track_id] = box
# 如果是右键单击事件
elif event == cv2.EVENT_RBUTTONDOWN:
# 清空已选择的盒子字典
self.selected_boxes = {}
# 重置左键点击计数为 0
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.
"""
# 获取第一个轨迹的盒子坐标并转换为 CPU 上的数组
self.boxes = tracks[0].boxes.xyxy.cpu()
# 获取第一个轨迹的类别并转换为 CPU 上的列表
self.clss = tracks[0].boxes.cls.cpu().tolist()
# 获取第一个轨迹的 ID 并转换为 CPU 上的列表
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.
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.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.keys():
if trk_id == track_id:
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
)
# 清空质心列表
self.centroids = []
# 如果需要显示图像并且环境检查通过，则显示图像帧
if self.view_img and self.env_check:
self.display_frames()
# 返回处理后的图像帧
return im0
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)
# 等待用户按键操作，如果按下 'q' 键则退出函数
if cv2.waitKey(1) & 0xFF == ord("q"):
return
if __name__ == "__main__":
# 当该脚本作为主程序运行时执行以下代码块
names = {0: "person", 1: "car"} # 示例类别名称的字典，键为索引，值为类别名称
# 创建 DistanceCalculation 的实例，传入类别名称的字典作为参数
distance_calculation = DistanceCalculation(names)
.\yolov8\ultralytics\solutions\heatmap.py
# 导入必要的库和模块
from collections import defaultdict # 导入collections模块中的defaultdict类
import cv2 # 导入OpenCV库
import numpy as np # 导入NumPy库
# 导入自定义的工具函数和类
from ultralytics.utils.checks import check_imshow, check_requirements # 导入检查函数
from ultralytics.utils.plotting import Annotator # 导入标注类
# 检查并确保所需的第三方库安装正确
check_requirements("shapely>=2.0.0")
# 导入用于空间几何计算的shapely库中的特定类和函数
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,
names,
imw=0,
imh=0,
colormap=cv2.COLORMAP_JET,
heatmap_alpha=0.5,
view_img=False,
view_in_counts=True,
view_out_counts=True,
count_reg_pts=None,
count_txt_color=(0, 0, 0),
count_bg_color=(255, 255, 255),
count_reg_color=(255, 0, 255),
region_thickness=5,
line_dist_thresh=15,
line_thickness=2,
decay_factor=0.99,
shape="circle",
):
"""Initializes the heatmap class with default values for Visual, Image, track, count and heatmap parameters."""
# Visual information
self.annotator = None # 初始化注释器为None
self.view_img = view_img # 设置是否显示图像
self.shape = shape # 设置热图形状
self.initialized = False # 标记对象是否已初始化
self.names = names # 类别名称列表
# Image information
self.imw = imw # 图像宽度
self.imh = imh # 图像高度
self.im0 = None # 初始化图像对象为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 # 初始化热图数组为None
self.heatmap_alpha = heatmap_alpha # 热图透明度
# Predict/track information
self.boxes = [] # 目标框列表
self.track_ids = [] # 跟踪目标的ID列表
self.clss = [] # 目标类别列表
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 = [] # 计数的目标ID列表
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) # 检查环境是否支持imshow函数
# Region and line selection
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.counting_region = LineString(self.count_reg_pts) # 使用两点创建线计数器区域
elif len(self.count_reg_pts) >= 3:
print("Polygon Counter Initiated.") # 打印多边形计数器初始化信息
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) # 使用线计数器作为默认选择
# 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):
"""
Extracts results from the provided data.
Args:
tracks (list): List of tracks obtained from the object tracking process.
"""
if tracks[0].boxes.id is not None:
self.boxes = tracks[0].boxes.xyxy.cpu() # 提取目标框坐标并转换为CPU格式
self.clss = tracks[0].boxes.cls.tolist() # 提取目标类别并转换为列表格式
self.track_ids = tracks[0].boxes.id.int().tolist() # 提取目标ID并转换为整数列表格式
def display_frames(self):
"""Display frames method."""
# 使用OpenCV显示图像窗口，标题为"Ultralytics Heatmap"，显示self.im0图像
cv2.imshow("Ultralytics Heatmap", self.im0)
# 等待用户按键输入，等待时间为1毫秒，并检查是否按下键盘上的q键
if cv2.waitKey(1) & 0xFF == ord("q"):
# 如果检测到按下q键，返回退出方法
return
if __name__ == "__main__":
# 如果当前脚本作为主程序执行
classes_names = {0: "person", 1: "car"} # 示例类别名称字典，映射类别编号到类别名称
# 创建一个 Heatmap 对象，传入类别名称字典作为参数
heatmap = Heatmap(classes_names)
.\yolov8\ultralytics\solutions\object_counter.py
# Ultralytics YOLO 🚀, AGPL-3.0 license
# 导入必要的库
from collections import defaultdict
import cv2
from ultralytics.utils.checks import check_imshow, check_requirements
from ultralytics.utils.plotting import Annotator, colors
# 检查并确保安装了必需的第三方库
check_requirements("shapely>=2.0.0")
# 导入 shapely 库中的几何图形类
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,
names,
reg_pts=None,
count_reg_color=(255, 0, 255),
count_txt_color=(0, 0, 0),
count_bg_color=(255, 255, 255),
line_thickness=2,
track_thickness=2,
view_img=False,
view_in_counts=True,
view_out_counts=True,
draw_tracks=False,
track_color=None,
region_thickness=5,
line_dist_thresh=15,
cls_txtdisplay_gap=50,
):
# 初始化对象计数器的各种参数
# names: 物体类别的名称列表
# reg_pts: 计数区域的定义点列表
# count_reg_color: 计数区域的颜色
# count_txt_color: 计数文本的颜色
# count_bg_color: 计数文本的背景颜色
# line_thickness: 绘制线条的粗细
# track_thickness: 绘制轨迹的粗细
# view_img: 是否显示图像
# view_in_counts: 是否显示进入计数区域的物体计数
# view_out_counts: 是否显示离开计数区域的物体计数
# draw_tracks: 是否绘制物体轨迹
# track_color: 轨迹颜色
# region_thickness: 计数区域的线条粗细
# line_dist_thresh: 线段连接的最大距离阈值
# cls_txtdisplay_gap: 不同类别文本显示的间隔
def mouse_event_for_region(self, event, x, y, flags, params):
"""
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 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):
if (
isinstance(point, (tuple, list))
and len(point) >= 2
and (abs(x - point[0]) < 10 and abs(y - point[1]) < 10)
):
self.selected_point = i
self.is_drawing = True
break
elif event == cv2.EVENT_MOUSEMOVE:
# 处理鼠标移动事件，如果正在绘制且选中了点，则更新计数区域的定义点
if self.is_drawing and self.selected_point is not None:
self.reg_pts[self.selected_point] = (x, y)
self.counting_region = Polygon(self.reg_pts)
elif event == cv2.EVENT_LBUTTONUP:
# 处理鼠标左键松开事件，停止绘制计数区域
self.is_drawing = False
self.selected_point = None
def display_frames(self):
"""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: # 如果用户绘制了计数区域，则添加鼠标事件处理
cv2.setMouseCallback(self.window_name, self.mouse_event_for_region, {"region_points": self.reg_pts})
cv2.imshow(self.window_name, self.im0)
# 检测按键事件，如果按下 'q' 键则关闭窗口
if cv2.waitKey(1) & 0xFF == ord("q"):
return
# 开始对象计数的主要函数，用于启动对象计数过程。
# 将当前帧从视频流存储到 self.im0 中
self.im0 = im0 # store image
# 对从对象跟踪过程获取的轨迹进行提取和处理
self.extract_and_process_tracks(tracks) # draw region even if no objects
# 如果 self.view_img 为 True，则显示帧
if self.view_img:
self.display_frames()
# 返回处理后的帧 self.im0
return self.im0
# 如果当前模块被直接运行（而不是被导入到其他模块中），则执行以下代码块
if __name__ == "__main__":
# 定义一个示例的类名字典，用于对象计数器
classes_names = {0: "person", 1: "car"} # example class names
# 创建一个对象计数器实例，传入类名字典作为参数
ObjectCounter(classes_names)
.\yolov8\ultralytics\solutions\parking_management.py
# Ultralytics YOLO 🚀, AGPL-3.0 license
import json # 导入处理 JSON 格式数据的模块
import cv2 # 导入 OpenCV 库，用于图像处理
import numpy as np # 导入 NumPy 库，用于处理数值数据
from ultralytics.utils.checks import check_imshow, check_requirements # 导入检查函数，用于检查必要的依赖项
from ultralytics.utils.plotting import Annotator # 导入绘图类，用于标注图像
class ParkingPtsSelection:
"""Class for selecting and managing parking zone points on images using a Tkinter-based UI."""
def __init__(self):
"""Initializes the UI for selecting parking zone points in a tkinter window."""
check_requirements("tkinter") # 检查是否安装了 tkinter 库，必要时抛出异常
import tkinter as tk # 导入 tkinter 库，用于构建图形用户界面
self.tk = tk # 赋值 tkinter 模块给实例变量 self.tk
self.master = tk.Tk() # 创建主窗口实例
self.master.title("Ultralytics Parking Zones Points Selector") # 设置窗口标题
# Disable window resizing
self.master.resizable(False, False) # 禁止窗口大小调整
# Setup canvas for image display
self.canvas = self.tk.Canvas(self.master, bg="white") # 在主窗口中创建画布用于显示图像
# Setup buttons
button_frame = self.tk.Frame(self.master) # 创建按钮框架
button_frame.pack(side=self.tk.TOP) # 放置在顶部
self.tk.Button(button_frame, text="Upload Image", command=self.upload_image).grid(row=0, column=0)
# 创建上传图像的按钮，点击后调用 upload_image 方法，放置在第一行第一列
self.tk.Button(button_frame, text="Remove Last BBox", command=self.remove_last_bounding_box).grid(
row=0, column=1
)
# 创建移除最后一个边界框的按钮，点击后调用 remove_last_bounding_box 方法，放置在第一行第二列
self.tk.Button(button_frame, text="Save", command=self.save_to_json).grid(row=0, column=2)
# 创建保存按钮，点击后调用 save_to_json 方法，放置在第一行第三列
# Initialize properties
self.image_path = None # 初始化图像路径为空
self.image = None # 初始化图像对象为空
self.canvas_image = None # 初始化画布图像对象为空
self.bounding_boxes = [] # 初始化边界框列表为空
self.current_box = [] # 初始化当前边界框为空
self.img_width = 0 # 初始化图像宽度为 0
self.img_height = 0 # 初始化图像高度为 0
# Constants
self.canvas_max_width = 1280 # 设置画布最大宽度为 1280
self.canvas_max_height = 720 # 设置画布最大高度为 720
self.master.mainloop() # 进入主事件循环，等待用户交互
def upload_image(self):
"""Upload an image and resize it to fit canvas."""
# 导入文件对话框模块
from tkinter import filedialog
# 导入PIL图像处理库及其图像展示模块ImageTk，因为ImageTk需要tkinter库
from PIL import Image, ImageTk
# 请求用户选择图片文件路径，限定文件类型为png、jpg、jpeg
self.image_path = filedialog.askopenfilename(filetypes=[("Image Files", "*.png;*.jpg;*.jpeg")])
if not self.image_path:
return # 如果未选择文件，则结束函数
# 打开选择的图片文件
self.image = Image.open(self.image_path)
self.img_width, self.img_height = self.image.size
# 计算图片的宽高比并调整图片大小以适应画布
aspect_ratio = self.img_width / self.img_height
if aspect_ratio > 1:
# 横向图片
canvas_width = min(self.canvas_max_width, self.img_width)
canvas_height = int(canvas_width / aspect_ratio)
else:
# 纵向图片
canvas_height = min(self.canvas_max_height, self.img_height)
canvas_width = int(canvas_height * aspect_ratio)
# 如果画布已经初始化，则销毁之前的画布对象
if self.canvas:
self.canvas.destroy()
# 创建新的画布对象，并设置其大小及背景色
self.canvas = self.tk.Canvas(self.master, bg="white", width=canvas_width, height=canvas_height)
# 调整图片大小，并转换为ImageTk.PhotoImage格式以在画布上展示
resized_image = self.image.resize((canvas_width, canvas_height), Image.LANCZOS)
self.canvas_image = ImageTk.PhotoImage(resized_image)
# 在画布上创建图片对象
self.canvas.create_image(0, 0, anchor=self.tk.NW, image=self.canvas_image)
# 将画布放置在窗口底部
self.canvas.pack(side=self.tk.BOTTOM)
# 绑定画布的鼠标左键点击事件到特定处理函数
self.canvas.bind("<Button-1>", self.on_canvas_click)
# 重置边界框和当前边界框数据
self.bounding_boxes = []
self.current_box = []
def on_canvas_click(self, event):
"""Handle mouse clicks on canvas to create points for bounding boxes."""
# 在画布上处理鼠标左键点击事件，用于创建边界框的顶点
self.current_box.append((event.x, event.y))
x0, y0 = event.x - 3, event.y - 3
x1, y1 = event.x + 3, event.y + 3
# 在画布上绘制红色的小圆点以标记边界框顶点
self.canvas.create_oval(x0, y0, x1, y1, fill="red")
if len(self.current_box) == 4:
# 如果当前边界框的顶点数为4，则将其添加到边界框列表中，并绘制边界框
self.bounding_boxes.append(self.current_box)
self.draw_bounding_box(self.current_box)
self.current_box = []
def draw_bounding_box(self, box):
"""
Draw bounding box on canvas.
Args:
box (list): Bounding box data
"""
# 在画布上绘制边界框
for i in range(4):
x1, y1 = box[i]
x2, y2 = box[(i + 1) % 4]
self.canvas.create_line(x1, y1, x2, y2, fill="blue", width=2)
# 从画布中移除最后一个绘制的边界框
def remove_last_bounding_box(self):
"""Remove the last drawn bounding box from canvas."""
from tkinter import messagebox # 为了多环境兼容性而导入消息框
# 如果存在边界框
if self.bounding_boxes:
self.bounding_boxes.pop() # 移除最后一个边界框
self.canvas.delete("all") # 清空画布
self.canvas.create_image(0, 0, anchor=self.tk.NW, image=self.canvas_image) # 重新绘制图像
# 重新绘制所有边界框
for box in self.bounding_boxes:
self.draw_bounding_box(box)
messagebox.showinfo("Success", "Last bounding box removed.") # 显示成功消息
else:
messagebox.showwarning("Warning", "No bounding boxes to remove.") # 显示警告消息：没有边界框可移除
# 将按图像到画布大小比例重新缩放的边界框保存到 'bounding_boxes.json'
def save_to_json(self):
"""Saves rescaled bounding boxes to 'bounding_boxes.json' based on image-to-canvas size ratio."""
from tkinter import messagebox # 为了多环境兼容性而导入消息框
canvas_width, canvas_height = self.canvas.winfo_width(), self.canvas.winfo_height()
width_scaling_factor = self.img_width / canvas_width
height_scaling_factor = self.img_height / canvas_height
bounding_boxes_data = []
# 遍历所有边界框
for box in self.bounding_boxes:
rescaled_box = []
for x, y in box:
rescaled_x = int(x * width_scaling_factor)
rescaled_y = int(y * height_scaling_factor)
rescaled_box.append((rescaled_x, rescaled_y))
bounding_boxes_data.append({"points": rescaled_box})
# 将数据以缩进格式写入到 'bounding_boxes.json'
with open("bounding_boxes.json", "w") as f:
json.dump(bounding_boxes_data, f, indent=4)
messagebox.showinfo("Success", "Bounding boxes saved to bounding_boxes.json") # 显示成功消息
class ParkingManagement:
"""Manages parking occupancy and availability using YOLOv8 for real-time monitoring and visualization."""
def __init__(
self,
model_path,
txt_color=(0, 0, 0),
bg_color=(255, 255, 255),
occupied_region_color=(0, 255, 0),
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() # 载入YOLOv8模型
# Labels dictionary
self.labels_dict = {"Occupancy": 0, "Available": 0} # 初始化标签字典
# Visualization details
self.margin = margin # 文字显示的边距
self.bg_color = bg_color # 背景颜色设置
self.txt_color = txt_color # 文字颜色设置
self.occupied_region_color = occupied_region_color # 占用区域的颜色设置
self.available_region_color = available_region_color # 空闲区域的颜色设置
self.window_name = "Ultralytics YOLOv8 Parking Management System" # 窗口名称
# Check if environment supports imshow
self.env_check = check_imshow(warn=True) # 检查环境是否支持imshow函数
def load_model(self):
"""Load the Ultralytics YOLO model for inference and analytics."""
from ultralytics import YOLO
return YOLO(self.model_path) # 使用路径加载Ultralytics YOLO模型
@staticmethod
def parking_regions_extraction(json_file):
"""
Extract parking regions from json file.
Args:
json_file (str): file that have all parking slot points
"""
with open(json_file, "r") as f:
return json.load(f) # 从JSON文件中提取停车区域信息
def process_data(self, json_data, im0, boxes, clss):
"""
Process the model data for parking lot management.
Args:
json_data (str): json data for parking lot management
im0 (ndarray): inference image
boxes (list): bounding boxes data
clss (list): bounding boxes classes list
Returns:
filled_slots (int): total slots that are filled in parking lot
empty_slots (int): total slots that are available in parking lot
"""
# 创建一个Annotator对象，用于在图像上标注信息
annotator = Annotator(im0)
# 初始化空车位数为json_data的长度，已占用车位数为0
empty_slots, filled_slots = len(json_data), 0
# 遍历json_data中的每个区域
for region in json_data:
# 将区域的点坐标转换为numpy数组形式
points_array = np.array(region["points"], dtype=np.int32).reshape((-1, 1, 2))
# 初始化区域占用状态为False
region_occupied = False
# 遍历所有检测到的边界框及其类别
for box, cls in zip(boxes, clss):
# 计算边界框中心点的坐标
x_center = int((box[0] + box[2]) / 2)
y_center = int((box[1] + box[3]) / 2)
# 获取类别名称对应的文本信息
text = f"{self.model.names[int(cls)]}"
# 在图像上显示对象标签信息
annotator.display_objects_labels(
im0, text, self.txt_color, self.bg_color, x_center, y_center, self.margin
)
# 计算当前中心点到区域边界的距离
dist = cv2.pointPolygonTest(points_array, (x_center, y_center), False)
# 如果距离大于等于0，表示中心点在区域内，标记该区域已被占用
if dist >= 0:
region_occupied = True
break
# 根据区域占用状态确定绘制区域的颜色
color = self.occupied_region_color if region_occupied else self.available_region_color
# 在图像上绘制多边形边界
cv2.polylines(im0, [points_array], isClosed=True, color=color, thickness=2)
# 如果区域被占用，更新已占用车位数和空车位数
if region_occupied:
filled_slots += 1
empty_slots -= 1
# 将已占用和空余车位数存入标签字典
self.labels_dict["Occupancy"] = filled_slots
self.labels_dict["Available"] = empty_slots
# 在图像上显示分析结果
annotator.display_analytics(im0, self.labels_dict, self.txt_color, self.bg_color, self.margin)
def display_frames(self, im0):
"""
Display frame.
Args:
im0 (ndarray): inference image
"""
# 如果开启了环境检测模式，创建并显示图像窗口
if self.env_check:
cv2.namedWindow(self.window_name)
cv2.imshow(self.window_name, im0)
# 检测键盘输入，如果按下 'q' 键，关闭窗口
if cv2.waitKey(1) & 0xFF == ord("q"):
return
.\yolov8\ultralytics\solutions\queue_management.py
# 引入 Python 内置的 collections 模块中的 defaultdict 类
from collections import defaultdict
# 引入 OpenCV 库
import cv2
# 引入自定义的检查函数
from ultralytics.utils.checks import check_imshow, check_requirements
# 引入自定义的绘图相关模块
from ultralytics.utils.plotting import Annotator, colors
# 检查运行环境是否满足要求，要求 shapely 版本 >= 2.0.0
check_requirements("shapely>=2.0.0")
# 引入 shapely 库中的几何对象 Point 和 Polygon
from shapely.geometry import Point, Polygon
class QueueManager:
"""A class to manage the queue in a real-time video stream based on object tracks."""
def __init__(
self,
names,
reg_pts=None,
line_thickness=2,
track_thickness=2,
view_img=False,
region_color=(255, 0, 255),
view_queue_counts=True,
draw_tracks=False,
count_txt_color=(255, 255, 255),
track_color=None,
region_thickness=5,
fontsize=0.7,
"""
Initializes the QueueManager with specified parameters for tracking and counting objects.
Args:
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 = names # 类别名称字典
self.annotator = None # 注释器对象
self.window_name = "Ultralytics YOLOv8 Queue Manager" # 窗口名称
# Object counting Information
self.counts = 0 # 对象计数初始为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.track_color = track_color # 设置轨迹线的颜色，如果为None则使用不同颜色区分不同轨迹
# Check if environment supports imshow
self.env_check = check_imshow(warn=True) # 检查环境是否支持imshow函数
def extract_and_process_tracks(self, tracks):
"""Extracts and processes tracks for queue management in a video stream."""
# 初始化注释器并绘制队列区域
self.annotator = Annotator(self.im0, self.tf, self.names)
# 检查是否有跟踪目标的盒子信息
if tracks[0].boxes.id is not None:
# 提取跟踪目标的盒子坐标并转换为CPU可处理的格式
boxes = tracks[0].boxes.xyxy.cpu()
# 提取类别信息并转换为列表
clss = tracks[0].boxes.cls.cpu().tolist()
# 提取跟踪目标的ID并转换为整数格式的列表
track_ids = tracks[0].boxes.id.int().cpu().tolist()
# 遍历每个跟踪目标
for box, track_id, cls in zip(boxes, track_ids, clss):
# 在图像上绘制边界框和标签
self.annotator.box_label(box, label=f"{self.names[cls]}#{track_id}", color=colors(int(track_id), True))
# 更新跟踪历史
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)
# 如果启用了绘制轨迹功能，则绘制轨迹
if self.draw_tracks:
self.annotator.draw_centroid_and_tracks(
track_line,
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
# 检查物体是否在计数区域内
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 = 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,
)
# 显示完成后重置计数
self.counts = 0
self.display_frames()
def display_frames(self):
"""Displays the current frame with annotations."""
if self.env_check and self.view_img:
# 绘制区域边界
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)
# 在按下 'q' 键时关闭窗口
if cv2.waitKey(1) & 0xFF == ord("q"):
return
# 存储当前帧到对象的实例变量中
self.im0 = im0 # Store the current frame
# 调用对象的方法，从跟踪列表中提取并处理跟踪信息
self.extract_and_process_tracks(tracks) # Extract and process tracks
# 如果视图图像标志为真，则显示当前帧
if self.view_img:
self.display_frames() # Display the frame if enabled
# 返回存储的当前帧
return self.im0
if __name__ == "__main__":
# 如果当前脚本作为主程序执行，则执行以下代码块
classes_names = {0: "person", 1: "car"} # 示例类别名称字典，将整数类别映射到字符串
queue_manager = QueueManager(classes_names)
# 创建一个队列管理器对象，使用给定的类别名称字典初始化
.\yolov8\ultralytics\solutions\speed_estimation.py
# Ultralytics YOLO 🚀, AGPL-3.0 license
from collections import defaultdict
from time import time
import cv2
import numpy as np
from ultralytics.utils.checks import check_imshow
from ultralytics.utils.plotting import Annotator, colors
class SpeedEstimator:
"""A class to estimate the speed of objects in a real-time video stream based on their tracks."""
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.
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 # 初始化原始图像为 None
self.annotator = None # 初始化标注器为 None
self.view_img = view_img # 设置是否显示图像的标志
# Region information
self.reg_pts = reg_pts if reg_pts is not None else [(20, 400), (1260, 400)] # 设置用于速度估计的区域点，默认为 [(20, 400), (1260, 400)]
self.region_thickness = region_thickness # 设置区域线的粗细
# Tracking information
self.clss = None # 初始化类别信息为 None
self.names = names # 设置类别名称字典
self.boxes = None # 初始化边界框信息为 None
self.trk_ids = None # 初始化跟踪 ID 信息为 None
self.line_thickness = line_thickness # 设置绘制框和轨迹线的粗细
self.trk_history = defaultdict(list) # 初始化跟踪历史为默认字典列表
# Speed estimation information
self.current_time = 0 # 初始化当前时间为 0
self.dist_data = {} # 初始化距离数据字典为空字典
self.trk_idslist = [] # 初始化跟踪 ID 列表为空列表
self.spdl_dist_thresh = spdl_dist_thresh # 设置速度计算的距离阈值
self.trk_previous_times = {} # 初始化上一个时间点的跟踪时间信息为空字典
self.trk_previous_points = {} # 初始化上一个时间点的跟踪点信息为空字典
# Check if the environment supports imshow
self.env_check = check_imshow(warn=True) # 检查环境是否支持 imshow 函数并设置警告为 True
def extract_tracks(self, tracks):
"""
Extracts results from the provided tracking data.
Args:
tracks (list): List of tracks obtained from the object tracking process.
"""
self.boxes = tracks[0].boxes.xyxy.cpu() # 提取边界框信息并转换为 CPU 格式
self.clss = tracks[0].boxes.cls.cpu().tolist() # 提取类别信息并转换为列表格式
self.trk_ids = tracks[0].boxes.id.int().cpu().tolist() # 提取跟踪 ID 并转换为整数列表格式
def store_track_info(self, track_id, box):
"""
存储跟踪数据。
Args:
track_id (int): 对象的跟踪ID。
box (list): 对象边界框数据。
Returns:
(list): 给定track_id的更新跟踪历史记录。
"""
# 获取当前跟踪ID对应的历史跟踪数据
track = self.trk_history[track_id]
# 计算边界框中心点坐标
bbox_center = (float((box[0] + box[2]) / 2), float((box[1] + box[3]) / 2))
# 将计算得到的中心点坐标添加到跟踪历史中
track.append(bbox_center)
# 如果跟踪历史长度超过30，移除最早的数据
if len(track) > 30:
track.pop(0)
# 将跟踪历史转换为numpy数组，并更新self.trk_pts
self.trk_pts = np.hstack(track).astype(np.int32).reshape((-1, 1, 2))
# 返回更新后的跟踪历史
return track
def plot_box_and_track(self, track_id, box, cls, track):
"""
绘制跟踪路径和边界框。
Args:
track_id (int): 对象的跟踪ID。
box (list): 对象边界框数据。
cls (str): 对象类别名称。
track (list): 用于绘制跟踪路径的跟踪历史。
"""
# 根据跟踪ID是否在速度数据中确定显示的速度标签
speed_label = f"{int(self.dist_data[track_id])} km/h" if track_id in self.dist_data else self.names[int(cls)]
# 根据跟踪ID是否在速度数据中确定绘制边界框的颜色
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):
"""
计算对象的速度。
Args:
trk_id (int): 对象的跟踪ID。
track (list): 用于绘制跟踪路径的跟踪历史。
"""
# 如果对象最后一个位置不在指定的区域内，则返回
if not self.reg_pts[0][0] < track[-1][0] < self.reg_pts[1][0]:
return
# 根据对象最后一个位置的y坐标是否在指定距离范围内确定运动方向
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"
# 如果前一次跟踪时间不为0，并且运动方向已知且跟踪ID不在列表中
if self.trk_previous_times.get(trk_id) != 0 and direction != "unknown" and trk_id not in self.trk_idslist:
# 将跟踪ID添加到列表中
self.trk_idslist.append(trk_id)
# 计算跟踪点的时间差和位置差，从而计算速度
time_difference = time() - self.trk_previous_times[trk_id]
if time_difference > 0:
dist_difference = np.abs(track[-1][1] - self.trk_previous_points[trk_id][1])
speed = dist_difference / time_difference
self.dist_data[trk_id] = speed
# 更新跟踪ID的前一次跟踪时间和位置
self.trk_previous_times[trk_id] = time()
self.trk_previous_points[trk_id] = track[-1]
def estimate_speed(self, im0, tracks, region_color=(255, 0, 0)):
"""
Estimates the speed of objects based on tracking data.
Args:
im0 (ndarray): Image.
tracks (list): List of tracks obtained from the object tracking process.
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=self.line_thickness)
# 绘制区域，使用给定的颜色和线条粗细
self.annotator.draw_region(reg_pts=self.reg_pts, color=region_color, thickness=self.region_thickness)
# 遍历每个框、轨迹ID和类别，并处理其信息
for box, trk_id, cls in zip(self.boxes, self.trk_ids, self.clss):
# 存储轨迹信息，并返回当前轨迹
track = self.store_track_info(trk_id, box)
# 如果当前轨迹ID不在之前时间的记录中，将其初始化为0
if trk_id not in self.trk_previous_times:
self.trk_previous_times[trk_id] = 0
# 绘制框和轨迹，并将其绘制到图像上
self.plot_box_and_track(trk_id, box, cls, track)
# 计算当前轨迹的速度
self.calculate_speed(trk_id, track)
# 如果视图模式开启且环境检查通过，显示当前帧图像
if self.view_img and self.env_check:
self.display_frames()
# 返回带有注释框和轨迹的图像
return im0
def display_frames(self):
"""Displays the current frame."""
# 显示当前帧图像，窗口标题为 "Ultralytics Speed Estimation"
cv2.imshow("Ultralytics Speed Estimation", self.im0)
# 检测键盘输入是否是 'q'，如果是则退出显示
if cv2.waitKey(1) & 0xFF == ord("q"):
return
if __name__ == "__main__":
# 如果这个脚本被直接执行而不是被导入为模块，则执行以下代码块
names = {0: "person", 1: "car"} # 示例类别名称，用于初始化速度估计器对象
speed_estimator = SpeedEstimator(names)
.\yolov8\ultralytics\solutions\streamlit_inference.py
# 导入所需的库
import io # 用于处理字节流
import time # 用于时间相关操作
import cv2 # OpenCV库，用于图像处理
import torch # PyTorch深度学习库
# 导入自定义函数和变量
from ultralytics.utils.checks import check_requirements # 导入检查依赖的函数
from ultralytics.utils.downloads import GITHUB_ASSETS_STEMS # 导入下载相关的变量
def inference(model=None):
"""使用Ultralytics YOLOv8在Streamlit应用中进行实时目标检测。"""
# 检查并确保Streamlit版本符合要求，以加快Ultralytics包的加载速度
check_requirements("streamlit>=1.29.0")
# 导入Streamlit库，仅在需要时进行导入以减少加载时间
import streamlit as st
# 导入YOLOv8模型
from ultralytics import YOLO
# 定义样式配置：隐藏主菜单
menu_style_cfg = """<style>MainMenu {visibility: hidden;}</style>"""
# 定义主标题配置：Ultralytics YOLOv8 Streamlit应用的标题
main_title_cfg = """<div><h1 style="color:#FF64DA; text-align:center; font-size:40px;
font-family: 'Archivo', sans-serif; margin-top:-50px;margin-bottom:20px;">
Ultralytics YOLOv8 Streamlit Application
</h1></div>"""
# 定义副标题配置：展示实时目标检测的描述
sub_title_cfg = """<div><h4 style="color:#042AFF; text-align:center;
font-family: 'Archivo', sans-serif; margin-top:-15px; margin-bottom:50px;">
Experience real-time object detection on your webcam with the power of Ultralytics YOLOv8! 🚀</h4>
</div>"""
# 设置Streamlit页面配置：页面标题、布局、侧边栏状态
st.set_page_config(page_title="Ultralytics Streamlit App", layout="wide", initial_sidebar_state="auto")
# 在页面中添加自定义的HTML样式和标题
st.markdown(menu_style_cfg, unsafe_allow_html=True)
st.markdown(main_title_cfg, unsafe_allow_html=True)
st.markdown(sub_title_cfg, unsafe_allow_html=True)
# 在侧边栏添加Ultralytics的Logo图标
with st.sidebar:
logo = "https://raw.githubusercontent.com/ultralytics/assets/main/logo/Ultralytics_Logotype_Original.svg"
st.image(logo, width=250)
# 在侧边栏添加标题：“用户配置”
st.sidebar.title("User Configuration")
# 添加视频源选择下拉菜单：webcam 或 video
source = st.sidebar.selectbox(
"Video",
("webcam", "video"),
)
vid_file_name = ""
if source == "video":
# 如果选择上传视频文件，则显示上传按钮
vid_file = st.sidebar.file_uploader("Upload Video File", type=["mp4", "mov", "avi", "mkv"])
if vid_file is not None:
g = io.BytesIO(vid_file.read()) # 将上传的视频文件读取为字节流对象
vid_location = "ultralytics.mp4"
with open(vid_location, "wb") as out: # 打开临时文件以写入字节
out.write(g.read()) # 将读取的字节写入文件
vid_file_name = "ultralytics.mp4"
elif source == "webcam":
vid_file_name = 0 # 如果选择使用摄像头，则设置视频源为默认摄像头
# 添加模型选择下拉菜单：从GITHUB_ASSETS_STEMS中选择以yolov8开头的模型
available_models = [x.replace("yolo", "YOLO") for x in GITHUB_ASSETS_STEMS if x.startswith("yolov8")]
if model:
available_models.insert(0, model.split(".pt")[0]) # 插入模型名称（去除.pt后缀）作为选项之一
selected_model = st.sidebar.selectbox("Model", available_models) # 选择所需的模型
with st.spinner("Model is downloading..."):
model = YOLO(f"{selected_model.lower()}.pt") # 加载 YOLO 模型
class_names = list(model.names.values()) # 将类名字典转换为类名列表
st.success("Model loaded successfully!") # 在界面上显示模型加载成功的消息
# 多选框，显示类名并获取所选类的索引
selected_classes = st.sidebar.multiselect("Classes", class_names, default=class_names[:3])
selected_ind = [class_names.index(option) for option in selected_classes]
if not isinstance(selected_ind, list): # 确保 selected_ind 是一个列表
selected_ind = list(selected_ind)
enable_trk = st.sidebar.radio("Enable Tracking", ("Yes", "No")) # 在侧边栏提供选择是否启用跟踪
conf = float(st.sidebar.slider("Confidence Threshold", 0.0, 1.0, 0.25, 0.01)) # 设置置信度阈值的滑块
iou = float(st.sidebar.slider("IoU Threshold", 0.0, 1.0, 0.45, 0.01)) # 设置IoU阈值的滑块
col1, col2 = st.columns(2)
org_frame = col1.empty() # 创建一个空白的列，用于显示原始帧
ann_frame = col2.empty() # 创建一个空白的列，用于显示带有注释的帧
fps_display = st.sidebar.empty() # 用于显示FPS的占位符
if st.sidebar.button("Start"): # 如果点击了“Start”按钮
videocapture = cv2.VideoCapture(vid_file_name) # 捕获视频
if not videocapture.isOpened():
st.error("Could not open webcam.") # 如果无法打开摄像头，则显示错误消息
stop_button = st.button("Stop") # 停止推断的按钮
while videocapture.isOpened():
success, frame = videocapture.read() # 读取视频帧
if not success:
st.warning("Failed to read frame from webcam. Please make sure the webcam is connected properly.")
break # 如果读取失败，则显示警告消息并退出循环
prev_time = time.time() # 记录当前时间
# 存储模型预测结果
if enable_trk == "Yes":
results = model.track(frame, conf=conf, iou=iou, classes=selected_ind, persist=True) # 调用模型进行跟踪
else:
results = model(frame, conf=conf, iou=iou, classes=selected_ind) # 调用模型进行推断
annotated_frame = results[0].plot() # 在帧上添加注释
# 计算模型的FPS
curr_time = time.time()
fps = 1 / (curr_time - prev_time)
prev_time = curr_time
# 在界面上显示原始帧和带注释的帧
org_frame.image(frame, channels="BGR")
ann_frame.image(annotated_frame, channels="BGR")
if stop_button:
videocapture.release() # 释放视频捕获资源
torch.cuda.empty_cache() # 清空CUDA内存
st.stop() # 停止Streamlit应用程序
# 在侧边栏显示FPS
fps_display.metric("FPS", f"{fps:.2f}")
videocapture.release() # 释放视频捕获资源
torch.cuda.empty_cache() # 清空CUDA内存
cv2.destroyAllWindows() # 销毁窗口
# 如果这个脚本被作为主程序执行（而不是被导入到其他脚本中），则执行以下代码
if __name__ == "__main__":
# 调用名为inference的函数来进行推断任务
inference()
.\yolov8\ultralytics\solutions\__init__.py
# 从当前包导入以下模块和函数
from .ai_gym import AIGym
from .analytics import Analytics
from .distance_calculation import DistanceCalculation
from .heatmap import Heatmap
from .object_counter import ObjectCounter
from .parking_management import ParkingManagement, ParkingPtsSelection
from .queue_management import QueueManager
from .speed_estimation import SpeedEstimator
from .streamlit_inference import inference
# 定义 __all__ 列表，用于指定当前模块导出的公共接口
__all__ = (
"AIGym", # 导出 AIGym 类
"DistanceCalculation", # 导出 DistanceCalculation 类
"Heatmap", # 导出 Heatmap 类
"ObjectCounter", # 导出 ObjectCounter 类
"ParkingManagement", # 导出 ParkingManagement 类
"ParkingPtsSelection", # 导出 ParkingPtsSelection 类
"QueueManager", # 导出 QueueManager 类
"SpeedEstimator", # 导出 SpeedEstimator 类
"Analytics", # 导出 Analytics 类
)
.\yolov8\ultralytics\trackers\basetrack.py
# 导入必要的模块和库
from collections import OrderedDict
import numpy as np
# 定义一个枚举类，表示被跟踪对象可能处于的不同状态
class TrackState:
"""
Enumeration class representing the possible states of an object being tracked.
Attributes:
New (int): State when the object is newly detected.
Tracked (int): State when the object is successfully tracked in subsequent frames.
Lost (int): State when the object is no longer tracked.
Removed (int): State when the object is removed from tracking.
"""
New = 0
Tracked = 1
Lost = 2
Removed = 3
# 定义对象跟踪的基类，提供基本属性和方法
class BaseTrack:
"""
Base class for object tracking, providing foundational attributes and methods.
Attributes:
_count (int): Class-level counter for unique track IDs.
track_id (int): Unique identifier for the track.
is_activated (bool): Flag indicating whether the track is currently active.
state (TrackState): Current state of the track.
history (OrderedDict): Ordered history of the track's states.
features (list): List of features extracted from the object for tracking.
curr_feature (any): The current feature of the object being tracked.
score (float): The confidence score of the tracking.
start_frame (int): The frame number where tracking started.
frame_id (int): The most recent frame ID processed by the track.
time_since_update (int): Frames passed since the last update.
location (tuple): The location of the object in the context of multi-camera tracking.
Methods:
end_frame: Returns the ID of the last frame where the object was tracked.
next_id: Increments and returns the next global track ID.
activate: Abstract method to activate the track.
predict: Abstract method to predict the next state of the track.
update: Abstract method to update the track with new data.
mark_lost: Marks the track as lost.
mark_removed: Marks the track as removed.
reset_id: Resets the global track ID counter.
"""
# 类级别的计数器，用于生成唯一的跟踪ID
_count = 0
def __init__(self):
"""Initializes a new track with unique ID and foundational tracking attributes."""
# 设置跟踪的唯一标识符
self.track_id = 0
# 标识当前跟踪是否处于激活状态
self.is_activated = False
# 设置跟踪的初始状态为新检测到
self.state = TrackState.New
# 记录跟踪状态的历史，使用有序字典维护状态顺序
self.history = OrderedDict()
# 存储从对象中提取的特征的列表
self.features = []
# 当前正在跟踪的对象特征
self.curr_feature = None
# 跟踪的置信度分数
self.score = 0
# 跟踪开始的帧编号
self.start_frame = 0
# 最近处理的帧编号
self.frame_id = 0
# 距离上次更新过去的帧数
self.time_since_update = 0
# 对象在多摄像头跟踪上下文中的位置
self.location = (np.inf, np.inf)
@property
def end_frame(self):
"""Return the last frame ID of the track."""
return self.frame_id
@staticmethod
def next_id():
"""Increment and return the global track ID counter."""
# 静态方法：递增并返回全局跟踪ID计数器
BaseTrack._count += 1
return BaseTrack._count
# 抽象方法：激活跟踪对象，接受任意数量的参数
def activate(self, *args):
"""Abstract method to activate the track with provided arguments."""
# 抛出未实现错误，子类需要实现具体的激活逻辑
raise NotImplementedError
# 抽象方法：预测跟踪对象的下一个状态
def predict(self):
"""Abstract method to predict the next state of the track."""
# 抛出未实现错误，子类需要实现具体的预测逻辑
raise NotImplementedError
# 抽象方法：使用新观测更新跟踪对象
def update(self, *args, **kwargs):
"""Abstract method to update the track with new observations."""
# 抛出未实现错误，子类需要实现具体的更新逻辑
raise NotImplementedError
# 将跟踪对象标记为丢失状态
def mark_lost(self):
"""Mark the track as lost."""
# 将跟踪对象状态设置为丢失
self.state = TrackState.Lost
# 将跟踪对象标记为移除状态
def mark_removed(self):
"""Mark the track as removed."""
# 将跟踪对象状态设置为移除
self.state = TrackState.Removed
# 静态方法：重置全局跟踪对象 ID 计数器
@staticmethod
def reset_id():
"""Reset the global track ID counter."""
# 将基类 BaseTrack 的跟踪对象计数器重置为 0
BaseTrack._count = 0
