记录发帖原因
今天是2024年12月27日,事情的起因是这样的,我们的人工智能老师要求我们做一个与人工智能相关的大作业,比如自然语言生成、语音识别、目标检测啥的课题,本人对目标检测这个课题有点兴趣,于是就选择了YOLO。(这篇文章可能篇幅较短,请见谅,本人这一周是考试周,后面有空还会再更新的)
1. 安装
关于YOLO是什么、谁做的、做来干嘛的,我就不多说了,想了解的可以去看其他人的介绍。
这是YOLO的官网:https://docs.ultralytics.com/quickstart/
我们根据教程的指引,直接输入指令安装即可:
pip install ultralytics
(后续还需要安装numpy
、opencv-python
等等之类的库)
2. 测试使用
安装完成之后,我们随便创建一个py文件,并在代码的同级目录下随便准备一张 demo.jpg 照片 ,把以下代码复制进去并运行:
from ultralytics import YOLO
if __name__ == '__main__':
model = YOLO("yolo11n.pt")
results = model("demo.jpg")
for result in results:
result.show()
(本人的这个demo.jpg照片是这个:
)
接下来我们运行代码:
发现正常运行,说明我们的YOLO安装没问题了。
上边这些一个个的矩阵框,我们在下文中都称其为盒子
。盒子
的左边是标签名称,右边的小数是对应该盒子的置信度
3. 训练
(小提醒:即使自己的笔记本电脑是游戏本,也不要轻易在游戏本上训练,不然显卡真会搞崩溃。。。推荐都在云算力机器上进行训练(应该都要钱))
训练py代码:
from ultralytics import YOLO
if __name__ == '__main__':
pre_pt_filename = "yolo11n.pt"
dataset_yaml_filename = "helmet_dataset"
model = YOLO(pre_pt_filename)
results = model.train(data=dataset_yaml_filename, epochs=100, imgsz=640, device=0, workers=18, batch=64)
其中,epochs
是训练轮次,workers
是训练的线程数量,batch
是批处理大小(即一次性训练处理的量)
dataset.yaml:
path: /root/work/fire
train: images/train
val: images/val
test: images/test
names:
0: fire
在 dataset.yaml
文件中,path
是数据集的总路径,后边的训练集train
、验证集val
、测试集test
都是相对于path
的相对路径。
以下是我在在线云算力机器上训练的过程截图 (在云算力机器上训练的话,可以尽情调整 batch
和 workers
来完全压榨算力机机能):
训练完成后,命令行会有提示。然后训练完成后的过程图和pt模型文件都会在 runs/detect/trainXXX 文件夹下(训练好的pt模型文件都在weights
文件夹下,一般来说,这里边会有两个pt文件,分别叫 best.pt
和 last.pt
。前者是所有epochs轮次训练的过程中,训练效果最好的一个,一般来说我们就用这个作为最终效果pt模型;后者则是保持有最后一轮epoch的pt模型,如果你想在训练了100次的基础上,再训练100次,就得用 last.pt
作为预训练模型再重新训练)
4. 摄像头检测
import cv2
from ultralytics import YOLO
model = YOLO("C:\\Users\\90965\\Desktop\\YOLO\\runs\\detect\\train\\weights\\best.pt")
video_path = 0
cap = cv2.VideoCapture(video_path)
while cap.isOpened():
success, frame = cap.read()
if success:
results = model(frame, conf=0.7)
for r in results:
for detection in r.boxes.xywh:
x = detection[0].item() # x坐标为列表中的第一个
y = detection[1].item() # y坐标为列表中的第二个
w = detection[2].item() # 矩形框宽度为列表中的第三个
h = detection[3].item() # 矩形框高度为列表中的第四个
print("中心点 x 坐标:", x)
print("中心点 y 坐标:", y)
annotated_frame = results[0].plot()
cv2.imshow('test', annotated_frame)
if cv2.waitKey(10) & 0xFF == ord("q"):
break
else:
break
cap.release()
cv2.destroyAllWindows()
上边的 model
填写为你所训练的best.pt文件的路径,这时候通过摄像头就可以看到你所训练的模型的效果。
conf
是最低置信度。假设设置此值为0.7
,则盒子的置信度若低于0.7,将会被直接忽略,只保留置信度高于0.7的盒子。
5. 获取盒子数据
model = YOLO("watch.pt")
video_path = 0
cap = cv2.VideoCapture(video_path)
while cap.isOpened():
success, frame = cap.read()
if success:
results = model(frame, conf=user_conf)
result_log = ""
num = 0
for r in results:
for i in range(len(r.boxes.xywh)):
x = r.boxes.xywh[i][0].item()
y = r.boxes.xywh[i][1].item()
w = r.boxes.xywh[i][2].item()
label_name = r.names[r.boxes.cls[i].item()] # 标签名称
if (label_name == "head"):
num += 1
cv2.putText(frame, str(num), (int(x), int(y)), cv2.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA)
cv2.circle(frame, (int(x), int(y)), int(w), (0, 0, 255), 2)
point_str = "侦测到未佩戴头盔者: 【" + str(num) + "】"
result_log += point_str + "\n"
elif (label_name == "helmet"):
num += 1
cv2.putText(frame, str(num), (int(x), int(y)), cv2.FONT_HERSHEY_COMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
len(r.boxes.xywh)
是被检测到的盒子数量
x = r.boxes.xywh[i][0].item()
y = r.boxes.xywh[i][1].item()
w = r.boxes.xywh[i][2].item()
h = r.boxes.xywh[i][3].item()
这四者分别是盒子的x坐标,y坐标,宽度和高度(也就是你在上文第2点看到的那张图的那一堆盒子)
通过 label_name = r.names[r.boxes.cls[i].item()]
可以获取盒子的标签名称(比如上文第2点的那张图里边的person
、chair
、cup
、tie
)
通过这些数据,你就可以自行干你想干的事情啦。