YOLOv5 目标检测算法

简介：

   目标检测在生存中应用范畴非常广泛，列如：道路违规抓拍、未戴口罩识别、工地未佩戴安全帽抓拍、厨房出现老鼠检测。
还可以用在游戏辅助外挂。以枪战为例，在游戏过程中时刻检测有没有人头出现。当检测到目标人头，再调用鼠标把枪口焦点移动到人头的位置，实现爆头效果。
本案例对yolov5官方代码进行了缩减，留下精华的部分，并封装成类。
yolov5默认模型支持80种目标检测，详细范例在文章最后。假如需要检测其他的目标，就需要自己网络数据进行数据标注，再重新训练新的模型。
  

---

1.参数初始化

1. def __init__(self):
2.     # classes，模型训练的时候每种目标都有自己的一个标识，类型数字从0开始。这里不指定，默认全部检测。
3.     self.classes = None
4.     # 加载pt模型
5.     self.weights = 'yolov5s.pt'
6.     self.imgsz = [640, 640]
7.     # 置信度，检测目标小于这个值的将不会被识别出来
8.     self.conf_thres = 0.5
9.     self.iou_thres = 0.45
10.     # Load model
11.     self.device = select_device('')
12.     self.model = DetectMultiBackend(self.weights, device=self.device, dnn=False, fp16=False)
13.     self.stride, self.names, self.pt = self.model.stride, self.model.names, self.model.pt
14.     self.imgsz = check_img_size(self.imgsz, s=self.stride)
15.     self.model.warmup(imgsz=(1, 3, *self.imgsz))

复制代码

2.算法推理

1. def detect(self, im0s):
2.     # 将图片缩放到640的大小进行识别
3.     img = letterbox(im0s, new_shape=(640, 640), stride=self.stride, auto=self.pt)[0]
4.     # Convert
5.     img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
6.     img = np.ascontiguousarray(img)
8.     im = torch.from_numpy(img).to(self.device)
9.     im = im.half() if self.model.fp16 else im.float()  # uint8 to fp16/32
10.     im /= 255
11.     if len(im.shape) == 3:
12.         im = im[None]
14.     pred = self.model(im, augment=False, visualize=False)
15.     pred = non_max_suppression(pred, self.conf_thres, self.iou_thres, self.classes, False, max_det=300)
17.     dets = []
18.     for i, det in enumerate(pred):
19.         # im0 为原图
20.         im0 = im0s.copy()
21.         if len(det):
22.             # 检测是以640进行的，所以需要把比例放到原图一样；
23.             det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()
24.             for *xyxy, conf, cls in reversed(det):
25.                 c = int(cls) # 检测目标对应的名称
26.                 # xyxy 包含了目标的坐标
27.                 dets.append([int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3]), self.names[c], f'{conf:.2f}'])

复制代码

3.调用算法：

2. from detect import SmokeUtil
3. import cv2
4. su = SmokeUtil()
6. if __name__ == "__main__":
7.     #1.视频
8.     cap = cv2.VideoCapture('target.mp4')
9.     while True:
10.         success, frame = cap.read()
11.         if not success:
12.             break
13.         dets = su.detect(frame)
15.         if len(dets) > 0:
16.             for j in dets:
17.                 cv2.rectangle(frame, (j[0], j[1]), (j[2], j[3]), (11, 44, 55), 3)
19.         cv2.imshow('show', frame)
20.         if cv2.waitKey(1) == ord('q'):
21.             break
22.     cap.release()
23.     cv2.destroyAllWindows()
25.     # 2.图片
26.     frame = cv2.imread('1.jpg')
27.     dets = su.detect(frame) #所有识别处理的坐标、名称、置信度
28.     if len(dets) > 0:
29.         for j in dets:
30.             cv2.rectangle(frame, (j[0], j[1]), (j[2], j[3]), (0,255,0), 3)
31.     cv2.imwrite('2.jpg', frame)

复制代码

4.检测效果——以大象为检测目标：

   （原图）  

   （检测效果）  

5.支持检测范例

1. ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck',
2. 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench',
3. 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra','giraffe',
4. 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard',
5. 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
6. 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl',
7. 'banana', 'apple', 'sandwich', 'orange', 'broccoli','carrot', 'hot dog', 'pizza',
8. 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet',
9. 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven',
10. 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
11. 'hair drier', 'toothbrush']

复制代码

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。