基于缓存队列的带抽帧功能的VideoCapture类

置顶 gaobiaoli

于 2023-12-13 15:19:33 发布

阅读量99

点赞数

文章标签：缓存 python opencv 目标检测机器学习

本文链接：https://blog.csdn.net/lrc199596/article/details/134972351

版权

一、需求来源

实时视频流与本地视频数据的抽帧检测任务（来源于实际科研-工程任务）
抽帧（间隔帧）读取
相机的畸变校正
多视频的时间对齐（帧偏移）

二、优化思路

缓存池+多线程实现帧的预读取
在计算模型初始化及计算过程中，实现视频的预读取，减少I/O的阻塞时间
cap.grab()代替cap.read()减少中间帧的解码时间
只在关键帧使用read()方法进行帧解码
可选undistort进行相机畸变校正
可选initStep用于初始的帧偏移

三、代码实现

项目完整代码开源至https://github.com/gaobiaoli/FasterVideoCapture欢迎提Issue与PR

import cv2
import threading
import queue


class BaseVideoCapture:
    """带畸变校正和间隔读取的视频播放器"""
    def __init__(self, videoPath, initStep=0, mtx=None, dist=None, interval=1) -> None:
        self.videoPath = videoPath
        self.mtx = mtx
        self.dist = dist
        self.capture = cv2.VideoCapture(videoPath)
        self._capture_count = 0
        self.initStep = initStep
        self.skip(step=initStep)
        self.interval = interval
    
    def skip(self, step):
        """跳过step帧"""
        for _ in range(step):
            ret = self._grab()
            if not ret:
                return False
        return True

    def _grab(self):
        """跳帧"""
        ret = self.capture.grab()
        if ret:
            self._capture_count += 1
        return ret

    def _read(self):
        """跳帧+读帧"""
        ret, frame = self.capture.read()
        if ret:
            self._capture_count += 1
        return ret, frame

    def read(self, interval=None):
        """间隔interval帧读取"""
        if not interval:
            interval = self.interval

        ret = self.skip(interval - 1)
        if not ret:
            return False, None
        
        ret, frame = self._read()
        if not ret:
            return False, None
        
        if self.mtx is not None and self.dist is not None:
            frame = cv2.undistort(frame, self.mtx, self.dist)
        return ret, frame

    def release(self):
        self.capture.release()
        
    def count(self):
        return self._capture_count

class FasterVideoCapture(BaseVideoCapture):
    VIDEO_END_FLAG = -1

    def __init__(
        self, videoPath, initStep=0, mtx=None, dist=None, interval=1, buffer_size=5
    ) -> None:
        super().__init__(videoPath, initStep, mtx, dist)
        self.interval = interval
        self.buffer_size = buffer_size
        self._read_count = initStep
        self.frame_buffer = queue.Queue(maxsize=buffer_size)
        self.stop_event = threading.Event()
        self.read_thread = threading.Thread(target=self._preload_frames, daemon=True)
        self.read_thread.start()

    def _preload_frames(self):
        while not self.stop_event.is_set():
            success = self.skip(self.interval - 1)
            if success:
                ret, frame = self._read()
                if ret:
                    self.frame_buffer.put((self._capture_count, frame), block=True)
            if not success or not ret:
                self.frame_buffer.put(
                    (FasterVideoCapture.VIDEO_END_FLAG, None), block=True
                )

    def read(self):
        count, frame = self.frame_buffer.get(block=True)
        if count == FasterVideoCapture.VIDEO_END_FLAG:
            return False, None
        self._read_count += self.interval
        assert count == self._read_count
        return True, frame

    def release(self):
        self.stop_event.set()
        self.read_thread.join()
        self.cap.release()

    def count(self):
        return self._read_count