一、环境配置
PyCharm 2024.1.4 (Community Edition)
Windows 11.0
cuda 12.4(目前nvidia其实已经有12.6了,但是pytorch目前只支持到12.4)
python 3.11.9(同样python已经有3.12了但是仅测试版有支持,而且现在conda好像也不支持3.12?)
YOLO 8.2.71
二、使用预训练的YOLOv8模型对视频进行处理
注意以下所有的程序都可以长按q结束,不再文档中注释
1.目标检测
# 从ultralytics包中导入YOLO类,这个类封装了YOLO模型的所有功能
from ultralytics import YOLO
# 创建一个YOLO类的实例,加载名为'yolov8n.pt'的预训练模型。
# 'yolov8n.pt'是YOLOv8的一个轻量级版本,适用于需要快速但可能牺牲一些精度的场景。
model = YOLO('yolov8n.pt')
# 使用加载的模型对视频源(这里是计算机的默认摄像头,用0表示)进行目标检测。
result = model(
source=0, # 指定视频源为计算机的默认摄像头。
show=True, # 检测到的结果会实时显示在屏幕上。
conf=0.4, # 设置置信度阈值为0.4,即只显示置信度大于或等于0.4的检测结果。
)
2.带fps显示的实例分割
import time # 导入time模块,用于测量模型推理时间
import cv2 # 导入OpenCV库,用于视频捕获和图像显示
from ultralytics import YOLO # 从ultralytics包中导入YOLO类,用于加载和运行YOLO模型
# 加载YOLOv8模型,这里使用的是'yolov8n-seg.pt',一个支持分割的轻量级模型
model = YOLO('yolov8n-seg.pt')
# 初始化摄像头,参数0通常代表计算机的默认摄像头
cap = cv2.VideoCapture(0)
# 当摄像头成功打开时,循环读取视频帧
while cap.isOpened():
# 从摄像头读取一帧
success, frame = cap.read()
# 如果成功读取到帧
if success:
# 记录当前时间,用于计算模型推理时间
start = time.perf_counter()
# 使用YOLO模型对帧进行推理
results = model(frame)
# 记录推理结束时间
end = time.perf_counter()
# 计算并获取总的推理时间
total_time = end - start
# 计算并获取帧率(FPS)
fps = 1 / total_time
# 从结果中获取第一帧的标注图像(对于实时视频,通常每次只处理一帧)
annotated_frame = results[0].plot()
# 在标注图像上添加FPS信息
cv2.putText(
annotated_frame, # annotated_frame: 标注后的图像
f'FPS:{int(fps)}', # f'FPS:{int(fps)}': 要绘制的FPS信息字符串
(10, 30), # (10, 30): 文本左下角的坐标
cv2.FONT_HERSHEY_SIMPLEX, # cv2.FONT_HERSHEY_SIMPLEX: 字体类型
0.7, # 0.7: 字体缩放比例因子
(0, 255, 0), # (0, 255, 0): 文本颜色,绿色
2, # 2: 线条粗细
cv2.LINE_AA # cv2.LINE_AA: 使用抗锯齿线条绘制文本
)
# 显示标注后的图像
cv2.imshow('YOLOv8 Inference', annotated_frame)
# 如果按下'q'键,则退出循环
if cv2.waitKey(1) & 0xFF == ord('q'):
break
3.姿态估计
from ultralytics import YOLO
# 加载YOLOv8姿态估计模型文件
model = YOLO('yolov8m-pose.pt')
# 使用模型进行推理
results = model(
source=0, # 表示使用默认摄像头作为输入
show=True, # 表示在屏幕上显示结果
conf=0.3, # 表示置信度阈值设置为0.3
)
4.目标检测和跟踪
from ultralytics import YOLO
model = YOLO('yolov8m.pt')
# 使用模型的.track()方法进行目标跟踪和检测
results = model.track(
source=0, # 通常表示使用计算机的默认摄像头作为视频源
show=True, # 表示在屏幕上显示结果
tracker='bytetrack.yaml' # 指定了跟踪算法的配置文件,这里是ByteTrack的配置
)
5.生成热力图
import cv2
from ultralytics import YOLO
from ultralytics.solutions import heatmap
# 加载YOLOv8s模型
model = YOLO('yolov8s.pt')
# 初始化摄像头
cap = cv2.VideoCapture(0)
# 初始化热力图对象
heatmap_obj = heatmap.Heatmap(
names=model.names, # 检测到的对象类别名称列表,来自YOLO模型的输出
view_img=True # 是否在内部显示图像
)
# 当视频捕获设备(如摄像头)处于打开状态时,执行循环
while cap.isOpened():
# 从视频捕获设备读取一帧图像
# success是一个布尔值,表示是否成功读取了帧
# im0是读取到的帧图像(如果成功的话)
success, im0 = cap.read()
results = model.track(
im0,
persist=True, # 表示跟踪上一帧的对象
show=False # 表示模型在跟踪时不应该显示任何结果
)
# 生成热力图
frame = heatmap_obj.generate_heatmap(im0, tracks=results)
# 按'q'键退出循环
if cv2.waitKey(1) & 0xFF == ord('q'):
break
6.导入视频进行目标计数并导出
import cv2
from ultralytics import YOLO
from ultralytics.solutions import object_counter
# 加载YOLOv8s模型,模型文件为'yolov8s.pt'
model = YOLO('yolov8s.pt')
# 初始化摄像头,但这里实际上是通过VideoCapture读取视频文件'fish.mp4'
# 如果要读取实时摄像头,应传递摄像头ID(通常是0)
cap = cv2.VideoCapture('fish.mp4') # 这里我导入了一个1920p的视频
# 初始化ObjectCounter对象,用于对象计数
counter = object_counter.ObjectCounter(
names=model.names, # names=model.names表示使用YOLO模型检测到的对象类别名称
draw_tracks=True, # draw_tracks=True表示在结果图像上绘制跟踪轨迹
view_img=True, # view_img=True表示在内部显示处理后的图像
reg_pts=[(300, 10), (400, 10), (400, 470), (300, 470)] # 逐点画框,注意顺序,这是1920p的视频但是我后面resize成480p了,所以画小点
)
# 设置视频输出参数
fourcc = cv2.VideoWriter_fourcc(*'mp4v') # 定义编解码器
out = cv2.VideoWriter('output.mp4', fourcc, 30, (720, 480)) # 创建VideoWriter对象
# 当视频捕获设备(这里是视频文件)处于打开状态时,执行循环
while cap.isOpened():
# 从视频捕获设备(或文件)读取一帧图像
# success是一个布尔值,表示是否成功读取了帧
# frame是读取到的图像帧
success, frame = cap.read()
# 如果成功读取到帧,则继续处理
if success:
# 将帧图像调整为720x480的大小
reframe = cv2.resize(frame, (720, 480))
# 使用YOLO模型进行对象跟踪
tracks = model.track(
reframe,
persist=True, # persist=True表示跟踪上一帧的对象
show=False # show=False表示模型在跟踪时不应该显示任何结果(结果将由ObjectCounter处理)
)
# 使用ObjectCounter对象进行对象计数,并绘制跟踪轨迹
counter.start_counting(reframe, tracks)
out.write(reframe)
# 按'q'键退出循环
if cv2.waitKey(1) & 0xFF == ord('q'):
break
else:
# 如果没有成功读取帧(即视频结束),则退出循环
break
# 释放资源
cap.release()
out.release()
cv2.destroyAllWindows()
导入视频为【可爱的fish们♿️(合集)-哔哩哔哩】 https://b23.tv/FgJWuBd
导出视频为【fish目标计数-哔哩哔哩】 https://b23.tv/Ar06CtU
7.实例分割和目标跟踪
import cv2
from ultralytics import YOLO
from ultralytics.utils.plotting import Annotator, colors # 导入Annotator类用于绘制和colors函数用于生成颜色
# 加载YOLOv8n的分割模型
model = YOLO("yolov8n-seg.pt") # 初始化模型,加载预训练的yolov8n-seg.pt权重文件
names = model.model.names # 获取模型识别的类别名称
# 使用OpenCV打开视频捕捉设备(通常是摄像头)
cap = cv2.VideoCapture(0)
# 无限循环,直到视频帧为空或用户中断
while True:
ret, im0 = cap.read() # 读取一帧视频
if not ret: # 如果读取失败(例如,摄像头被关闭)
print("Video frame is empty or video processing has been successfully completed.")
break # 跳出循环
# 使用YOLO模型对图像进行预测
results = model.predict(im0) # results包含检测到的对象的边界框、类别、置信度等信息,对于分割模型还包括mask信息
# 初始化Annotator对象,用于在图像上绘制结果
annotator = Annotator(im0, line_width=2)
# 检查是否有mask信息(即是否有分割结果)
if results[0].masks is not None:
clss = results[0].boxes.cls.cpu().tolist() # 获取检测到的类别的索引(CPU版本)
masks = results[0].masks.xy # 获取分割的mask信息
# 遍历每个mask和对应的类别
for mask, cls in zip(masks, clss):
color = colors(int(cls), True) # 根据类别索引生成颜色
txt_color = annotator.get_txt_color(color) # 获取与mask颜色对比度高的文本颜色
# 在图像上绘制分割结果,包括mask和类别标签
annotator.seg_bbox(mask=mask, mask_color=color, label=names[int(cls)], txt_color=txt_color)
# 显示处理后的图像
cv2.imshow("instance-segmentation", im0)
# 如果按下'q'键,则退出循环
if cv2.waitKey(1) & 0xFF == ord("q"):
break
8.速度估计
import cv2
from ultralytics import YOLO, solutions
# 加载YOLOv8模型
model = YOLO("yolov8n.pt") # 使用预训练的 yolov8n 模型进行初始化
names = model.model.names # 获取模型可以识别的对象名称列表
# 打开摄像头
cap = cv2.VideoCapture(0) # 打开默认摄像头(通常是笔记本或计算机的内置摄像头)
# 定义一条用于速度估计的线条(假设视频分辨率为1280x720,360是y坐标)
line_pts = [(0, 360), (1280, 360)] # 线条从屏幕左侧到右侧,穿过预定义的高度
# 初始化速度估计对象
speed_obj = solutions.SpeedEstimator(
reg_pts=line_pts, # 线条点,用于计算通过该线条的对象的速度
names=names, # 对象名称列表,用于识别哪些对象的速度需要被估计
view_img=True, # 是否在结果图像上显示速度估计
)
# 主循环,处理视频流
while cap.isOpened():
success, im0 = cap.read() # 从摄像头读取一帧
# 使用YOLOv8进行对象跟踪,但不显示跟踪结果
tracks = model.track(im0, persist=True, show=False)
# 估计并显示速度
im0 = speed_obj.estimate_speed(im0, tracks)
# 如果按下'q'键,则退出循环
if cv2.waitKey(1) & 0xFF == ord("q"):
break
9.导入视频进行运动计数并导出
import cv2
from ultralytics import YOLO, solutions
# 加载预训练的YOLOv8n-pose模型,用于姿态估计
model = YOLO("yolov8n-pose.pt")
# 导入俯卧撑视频
cap = cv2.VideoCapture('pushup.mp4')
fourcc = cv2.VideoWriter_fourcc(*'mp4v') # 定义编解码器
out = cv2.VideoWriter('output.mp4', fourcc, 30, (1080, 1920)) # 创建VideoWriter对象
# 创建一个AIGym对象,用于基于姿态估计的结果进行计数或分析
# 参数包括线条粗细、是否显示图像、要检测的姿态类型(如俯卧撑)以及要检查的关键点索引
gym_object = solutions.AIGym(
line_thickness=2, # 绘制线条的粗细
view_img=True, # 是否在屏幕上显示处理后的图像
pose_type="pushup",
# pose_type 参数是一个可选的字符串(str),用于指定要检测的姿态类型。
# 这个参数可以接受三个预定义的值:'pullup'(引体向上)、'pushup'(俯卧撑)或 'abworkout'(腹部锻炼)。
# 如果调用时不提供这个参数,它将默认为 "pullup"。
kpts_to_check=[6, 8, 10], # 要检查的关键点索引,这些索引对应于特定的身体部位
)
# 使用视频捕获对象进行循环,直到视频流结束
while cap.isOpened():
# 从摄像头读取一帧图像
success, im0 = cap.read()
# 如果成功读取到图像,则使用YOLO模型进行跟踪(这里实际上是姿态估计)
# verbose=False表示不在控制台输出跟踪信息
results = model.track(im0, verbose=False) # 注意:YOLOv8通常用于检测,但这里假设它被用于跟踪或姿态估计
# 使用AIGym对象处理图像和结果,可能包括绘制线条、计数等
im0 = gym_object.start_counting(im0, results)
out.write(im0)
# 如果按下'q'键,则退出循环
if cv2.waitKey(1) & 0xFF == ord("q"):
break
# 释放资源
cap.release()
out.release()
cv2.destroyAllWindows()
10.模糊对象(马赛克)
import cv2
from ultralytics import YOLO
from ultralytics.utils.plotting import Annotator, colors
model = YOLO("yolov8n.pt")
# 获取模型可以检测到的对象名称
names = model.names
cap = cv2.VideoCapture(0)
# 设置模糊比例
blur_ratio = 50
# 当摄像头打开时,循环读取帧
while cap.isOpened():
# 从摄像头读取一帧
success, im0 = cap.read()
# 使用YOLO模型进行预测,不显示预测结果
results = model.predict(im0, show=False)
# 提取预测结果中的边界框和类别
boxes = results[0].boxes.xyxy.cpu().tolist() # 边界框坐标(x_min, y_min, x_max, y_max)
clss = results[0].boxes.cls.cpu().tolist() # 类别索引
# 创建一个Annotator对象用于在图像上绘制
annotator = Annotator(im0, line_width=2, example=names)
# 如果有边界框,则遍历它们
if boxes is not None:
for box, cls in zip(boxes, clss):
# 在图像上绘制边界框和标签
annotator.box_label(box, color=colors(int(cls), True), label=names[int(cls)])
# 裁剪出边界框内的对象
obj = im0[int(box[1]):int(box[3]), int(box[0]):int(box[2])]
# 对裁剪出的对象进行模糊处理
blur_obj = cv2.blur(obj, (blur_ratio, blur_ratio))
# 将模糊后的对象替换回原图中的相应位置
im0[int(box[1]):int(box[3]), int(box[0]):int(box[2])] = blur_obj
# 显示处理后的图像
cv2.imshow("ultralytics", im0)
# 如果按下'q'键,则退出循环
if cv2.waitKey(1) & 0xFF == ord("q"):
break
三、Gradio与YOLOv8集成
import gradio as gr
import PIL.Image as Image
from ultralytics import ASSETS, YOLO
model = YOLO("yolov8n.pt")
# 定义一个函数,用于预测图像中的对象
def predict_image(img, conf_threshold, iou_threshold):
"""
使用YOLOv8模型预测图像中的对象,允许调整置信度和IOU阈值。
参数:
img (PIL.Image): 要预测的PIL图像。
conf_threshold (float): 置信度阈值,低于此值的检测结果将被忽略。
iou_threshold (float): IOU(交并比)阈值,用于非极大值抑制。
返回:
PIL.Image: 包含检测结果的PIL图像。
"""
# 使用模型进行预测
results = model.predict(source=img, conf=conf_threshold, iou=iou_threshold, show_labels=True, show_conf=True,
imgsz=640)
# 遍历结果
for r in results:
# 绘制检测结果
im_array = r.plot()
# 将绘制结果的numpy数组转换为PIL图像。
im = Image.fromarray(im_array[..., ::-1])
return im
# 创建Gradio界面
iface = gr.Interface(fn=predict_image, # 指定函数
inputs=[ # 输入组件
gr.Image(type="pil", label="图像文件上传"), # PIL类型的图像输入
gr.Slider(minimum=0, maximum=1, value=0.25, label="置信度阈值"), # 置信度阈值滑块
gr.Slider(minimum=0, maximum=1, value=0.45, label="交并比(IoU)阈值"), # IOU阈值滑块
], outputs=gr.Image(type="pil", label="Result"), # 输出组件,PIL类型的图像
title="Ultralytics Gradio", # 界面标题
description="上传图像以供推理。默认情况下使用Ultralytics的YOLOv8n模型。",
# 界面描述
examples=[[ASSETS / "bus.jpg", 0.25, 0.45], [ASSETS / "zidane.jpg", 0.25, 0.45]],
)
# 启动Gradio界面
if __name__ == "__main__":
iface.launch()
结果如下
四、使用yolo对数据集训练图像分类模型
from ultralytics import YOLO
def main():
model = YOLO("yolov8n-cls.pt")
# 尝试使用CIFAR-10数据集进行训练
results = model.train(
# CIFAR-10数据集包含60,000张32x32像素的彩色(RGB)图像。这些图像被分为10个不同的类别,每个类别包含6,000张图像。
# 具体类别包括飞机(airplane)、汽车(automobile)、鸟类(bird)、猫(cat)、鹿(deer)、狗(dog)、蛙(frog)、马(horse)、船(ship)和卡车(truck)。
# 数据集被进一步划分为训练集和测试集。训练集包含50,000张图像,用于模型的训练过程;测试集包含10,000张图像,用于评估训练好的模型的性能。
data="cifar10",
epochs=50, # 训练50次
imgsz=32 # 设置输入图像的大小为32x32像素
)
if __name__ == '__main__':
# 这里调用 main() 函数作为程序的入口点
# 只有当文件被直接运行时,main() 函数才会被执行
# 提高模块的可重用性:通过将主程序逻辑封装在if __name__ == '__main__': 块中,你可以确保当模块被其他文件导入时,不会执行不必要的代码。
# 这使得你的模块更加干净、更加容易在其他上下文中重用。
main()
跑出来的结果差不多是这样,环境如下
Python-3.11.9 torch-2.4.0 CUDA:0 (NVIDIA GeForce RTX 4060 Laptop GPU, 8188MiB)