简介：

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

1.参数初始化

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

2.算法推理

def detect(self, im0s):
    # 将图片缩放到640的大小进行识别
    img = letterbox(im0s, new_shape=(640, 640), stride=self.stride, auto=self.pt)[0]
    # Convert
    img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
    img = np.ascontiguousarray(img)

    im = torch.from_numpy(img).to(self.device)
    im = im.half() if self.model.fp16 else im.float()  # uint8 to fp16/32
    im /= 255
    if len(im.shape) == 3:
        im = im[None]

    pred = self.model(im, augment=False, visualize=False)
    pred = non_max_suppression(pred, self.conf_thres, self.iou_thres, self.classes, False, max_det=300)

    dets = []
    for i, det in enumerate(pred):
        # im0 为原图
        im0 = im0s.copy()
        if len(det):
            # 检测是以640进行的，所以需要把比例放到原图一样；
            det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round()
            for *xyxy, conf, cls in reversed(det):
                c = int(cls) # 检测目标对应的名称
                # xyxy 包含了目标的坐标
                dets.append([int(xyxy[0]), int(xyxy[1]), int(xyxy[2]), int(xyxy[3]), self.names[c], f'{conf:.2f}'])

3.调用算法：

from detect import SmokeUtil
import cv2
su = SmokeUtil()

if __name__ == "__main__":
    #1.视频
    cap = cv2.VideoCapture('target.mp4')
    while True:
        success, frame = cap.read()
        if not success:
            break
        dets = su.detect(frame)

        if len(dets) > 0:
            for j in dets:
                cv2.rectangle(frame, (j[0], j[1]), (j[2], j[3]), (11, 44, 55), 3)

        cv2.imshow('show', frame)
        if cv2.waitKey(1) == ord('q'):
            break
    cap.release()
    cv2.destroyAllWindows()

    # 2.图片
    frame = cv2.imread('1.jpg')
    dets = su.detect(frame) #所有识别处理的坐标、名称、置信度
    if len(dets) > 0:
        for j in dets:
            cv2.rectangle(frame, (j[0], j[1]), (j[2], j[3]), (0,255,0), 3)
    cv2.imwrite('2.jpg', frame)

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

5.支持检测类型

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