import multiprocessing import cv2 import threading import os class VideoReader: listOfFrames = None w = None h = None def __init__(self, config, setOfFrames=None): videoPath = config["inputPath"] if videoPath is None: raise Exception("ERROR: Video reader needs a videoPath!") self.videoPath = videoPath self.lastFrame = 0 # buffer data struct: # buffer = Queue([(frameNumber, frame), ]) self.buffer = multiprocessing.Queue(config["videoBufferLength"]) #self.vc = cv2.VideoCapture(videoPath) self.stopped = False self.getWH() self.calcFPS() self.calcLength() self.calcStartTime() if setOfFrames is not None: self.listOfFrames = sorted(setOfFrames) def __enter__(self): self.fillBuffer() return self def __exit__(self, type, value, traceback): self.stop() def stop(self): self.thread.join() self.vc.release() def pop(self): frameNumber, frame = self.buffer.get(block=True) if frame is None: self.stopped = True return frameNumber, frame def fillBuffer(self, listOfFrames=None): self.endFrame = int(cv2.VideoCapture(self.videoPath).get(cv2.CAP_PROP_FRAME_COUNT)) if listOfFrames is not None: self.listOfFrames = listOfFrames if self.listOfFrames is not None: self.thread = multiprocessing.Process(target=self.readFramesByList, args=()) else: self.thread = multiprocessing.Process(target=self.readFrames, args=()) self.thread.start() def readFrames(self): """Reads video from start to finish""" self.vc = cv2.VideoCapture(self.videoPath) while self.lastFrame < self.endFrame: res, frame = self.vc.read() if res: self.buffer.put((self.lastFrame, frame)) self.lastFrame += 1 self.buffer.put((self.lastFrame, None)) def readFramesByList(self): """Reads all frames from a list of frame numbers""" self.vc = cv2.VideoCapture(self.videoPath) self.vc.set(1, self.listOfFrames[0]) self.lastFrame = self.listOfFrames[0] self.endFrame = self.listOfFrames[-1] while self.lastFrame < self.endFrame: if self.lastFrame in self.listOfFrames: res, frame = self.vc.read() if res: self.buffer.put((self.lastFrame, frame)) else: print("Couldn't read Frame") # since the list is sorted the first element is always the lowest relevant framenumber # [0,1,2,3,32,33,34,35,67,68,69] self.listOfFrames.pop(0) self.lastFrame += 1 else: # if current Frame number is not in list of Frames, we can skip a few frames self.vc.set(1, self.listOfFrames[0]) self.lastFrame = self.listOfFrames[0] self.buffer.put((self.lastFrame, None)) def videoEnded(self): if self.stopped and self.buffer.empty(): return True else: return False def calcFPS(self): self.fps = cv2.VideoCapture(self.videoPath).get(cv2.CAP_PROP_FPS) def getFPS(self): if self.fps is None: self.calcFPS() return self.fps def calcLength(self): fc = int(cv2.VideoCapture(self.videoPath).get(cv2.CAP_PROP_FRAME_COUNT)) self.length = fc / self.getFPS() def getLength(self): if self.length is None: self.calcLength() return self.length def calcStartTime(self): starttime = os.stat(self.videoPath).st_mtime length = self.getLength() starttime = starttime - length self.starttime = starttime def getStartTime(self): return self.starttime def getWH(self): """get width and height""" vc = cv2.VideoCapture(self.videoPath) if self.w is None or self.h is None: res, image = vc.read() self.w = image.shape[1] self.h = image.shape[0] return (self.w, self.h)