30行代码让你轻松玩转Yolov5

30行代码让你轻松玩转Yolov5

最近，有小伙伴私信我说在做项目的时候不需要那么多Yolo的文件，问能不能出一个简化的Yolo教程，今天，他来了。

首先要导入库包

import cv2
import torch
import time
import numpy as np

cv2是opencv的安装包，如果报错：

No module named '×××'（其中xxx表示包的名称)

那么解决办法就是

pip install xxx(特别地，cv2叫opencv-python,sklearn叫scikit-learn)，即：
缺少cv2:pip install opencv-python
缺少sklearn:pip install scikit-learn

加载模型和设置模型参数

model = torch.hub.load('../yolov5', 'custom', path='../yolov5/weights/yolov5n.pt', source='local',device='cpu')
f = 0.4

要想利用cuda推理，就把上面的device='cpu'改为device='cuda0'即可。

要用到摄像头肯定得用到下面这行代码：

cap = cv2.VideoCapture(0)

0表示本机上可用的第0个摄像头，在没有接外接摄像头的情况下，0就是电脑摄像头。当接到外接摄像头时，可以用该外接摄像头的编号。此外，还可以接远程摄像头，远程摄像头将代码推流到服务器，cv2.VideoCapture还可以传入视频连接参数作为检测对象。详细内容见我之前写的食物分类移动端部署：图像分类之食物分类（NCWU机器学习课程设计）

之后就可以对传来的图像进行处理了，代码如下：

while True:ret,frame = ad()frame = cv2.flip(frame,1)img_cvt = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)results = model(img_cvt)print(str(results.pandas().xyxy[0].to_numpy()[:,-1]))# tensor-to-numpyresults_ = results.pandas().xyxy[0].to_numpy()i = 0

上述代码是读取摄像头里面的内容，返回一个ret代表是否成功读取，frame就是读取到的图像。当ret为False时，frame为None。之后将读取到的图像转为RGB类型，因为opencv读取的时BGR格式的图像。将读到的图像经过模型推理按得到结果results。

接着在图像上画定位框，代码如下

      for box in results_:l,t,r,b = box[:4].astype('int')confidence = str(round(box[4]*100,2))+"%"cls_name = box[angle(frame,(l,t),(r,b),(0,255,0),2)cv2.putText(frame,cls_name + "-" + confidence,(l,t),cv2.FONT_ITALIC,1,(255,0,0),2)cv2.imshow("result",frame)if cv2.waitKey(10) & 0xFF == ord("q"):break

最后的结果如下：

相比​ultralytics大神​的代码，我这个代码算是简单到令人发指。总结一下，全文的代码一共28行，其中还有几行是没用的，各位大佬在用的时候可以根据自己的需求对我的代码进行魔改。完整代码：

import cv2
import torch
import time
import numpy as npmodel = torch.hub.load('../yolov5', 'custom', path='../yolov5/weights/yolov5n.pt', source='local',device='cpu')
f = 0.4
cap = cv2.VideoCapture(0)
while True:ret,frame = ad()frame = cv2.flip(frame,1)img_cvt = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)results = model(img_cvt)print(str(results.pandas().xyxy[0].to_numpy()[:,-1]))# tensor-to-numpyresults_ = results.pandas().xyxy[0].to_numpy()i = 0for box in results_:l,t,r,b = box[:4].astype('int')confidence = str(round(box[4]*100,2))+"%"cls_name = box[angle(frame,(l,t),(r,b),(0,255,0),2)cv2.putText(frame,cls_name + "-" + confidence,(l,t),cv2.FONT_ITALIC,1,(255,0,0),2)cv2.imshow("result",frame)if cv2.waitKey(10) & 0xFF == ord("q"):break
lease()
cv2.destroyAllWindows()