目标检测——使用yolov6进行视频推理

yolov6源码里自带了图片的推理，推理完成后结果图片保存再runs/inference文件夹下；

但当我们想对视频进行推理的时候不是很方便，但v6里面封装了对图片的推理函数，我们只需要稍微改一下即可。

默认的图片推理代码如下：

    inferer = Inferer(source, weights, device, yaml, img_size, half)
    inferer.infer(conf_thres, iou_thres, classes, agnostic_nms, max_det, save_dir, save_txt, save_img, hide_labels, hide_conf)

我们只需要基于图片的infer代码，修改出一个infer_video接口即可：

找到inferer.py，定义两个新函数，其实就是复制infer和process_image函数并略微修改

1. infer_video

    def infer_video(self, cap, conf_thres, iou_thres, classes, agnostic_nms, max_det, save_dir, save_txt, save_img, hide_labels, hide_conf):
        cv2.namedWindow("test", cv2.WINDOW_NORMAL)
        while True:
            r, frame = cap.read()
            if not r:
                break
            else:
                img, img_src = self.precess_video(frame, self.img_size, self.stride, self.half)
                img = img.to(self.device)
                if len(img.shape) == 3:
                    img = img[None]
                    # expand for batch dim
                pred_results = self.model(img)
                det = non_max_suppression(pred_results, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)[0]

                gn = torch.tensor(img_src.shape)[[1, 0, 1, 0]]  # normalization gain whwh
                img_ori = img_src

                # check image and font
                assert img_ori.data.contiguous, 'Image needs to be contiguous. Please apply to input images with np.ascontiguousarray(im).'
                self.font_check()

                if len(det):
                    det[:, :4] = self.rescale(img.shape[2:], det[:, :4], img_src.shape).round()
                    for *xyxy, conf, cls in reversed(det):
                        class_num = int(cls)  # integer class
                        label = None if hide_labels else (self.class_names[
                                                              class_num] if hide_conf else f'{self.class_names[class_num]} {conf:.2f}')

                        self.plot_box_and_label(img_ori, max(round(sum(img_ori.shape) / 2 * 0.003), 2), xyxy, label,
                                                color=self.generate_colors(class_num, True))

                img_src = np.asarray(img_ori)
                cv2.imshow("test", img_src)
                if cv2.waitKey(33) == ord("q"):
                    break

2. precess_video

@staticmethod
    def precess_video(img, img_size, stride, half):
        '''Process image before image inference.'''
        try:
            img_src = img
            assert img_src is not None, f'Invalid image: {img}'
        except Exception as e:
            LOGGER.Warning(e)
        image = letterbox(img_src, img_size, stride=stride)[0]

        # Convert
        image = image.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
        image = torch.from_numpy(np.ascontiguousarray(image))
        image = image.half() if half else image.float()  # uint8 to fp16/32
        image /= 255  # 0 - 255 to 0.0 - 1.0

        return image, img_src

然后回到infer.py，我们把推理图片的代码注掉，使用下面代码进行推理即可

# Inference
    inferer = Inferer(source, weights, device, yaml, img_size, half)
    # infer imgs
    #inferer.infer(conf_thres, iou_thres, classes, agnostic_nms, max_det, save_dir, save_txt, save_img, hide_labels, hide_conf)
    # infer video
    videolist = ["../data/videos/test.MOV",
                 "../data/videos/test1.MOV",
                 "../data/videos/test2.mp4",
                 "../data/videos/test3.mp4",
                 "../data/videos/genvideo.avi"]
    cap = cv2.VideoCapture(videolist[4])
    inferer.infer_video(cap, conf_thres, iou_thres, classes, agnostic_nms, max_det, save_dir, save_txt, save_img, hide_labels, hide_conf)