Object detection - VIDEO reasoning using yolov6

yolov6 The source code comes with the reasoning of pictures , After the reasoning is completed, the result picture is saved and then runs/inference Under the folder ;

But it is not very convenient when we want to reason about video , but v6 It encapsulates the inference function of the picture , We just need to change it a little .

The default image reasoning code is as follows ：

    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)

We just need a picture based infer Code , Modify a infer_video Interface can ：

find inferer.py, Define two new functions , It's actually copying infer and process_image Function and slightly modify

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

Then back infer.py, Let's note out the code of the reasoning picture , Use the following code for reasoning

# 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)