added context manager for video reader

This commit is contained in:
Askill 2020-12-29 10:41:55 +01:00
parent 8627608d1e
commit e73fafcb9f
2 changed files with 35 additions and 16 deletions

View File

@ -1,6 +1,3 @@
from Application.Config import Config
from datetime import datetime
from queue import Queue from queue import Queue
import cv2 import cv2
@ -9,14 +6,13 @@ import os
class VideoReader: class VideoReader:
listOfFrames = None listOfFrames = None
w = 0 w = None
h = 0 h = None
def __init__(self, config, setOfFrames=None): def __init__(self, config, setOfFrames=None):
videoPath = config["inputPath"] videoPath = config["inputPath"]
if videoPath is None: if videoPath is None:
raise Exception("ERROR: Video reader needs a videoPath!") raise Exception("ERROR: Video reader needs a videoPath!")
self.videoPath = videoPath self.videoPath = videoPath
self.lastFrame = 0 self.lastFrame = 0
# buffer data struct: # buffer data struct:
@ -30,13 +26,17 @@ class VideoReader:
self.calcStartTime() self.calcStartTime()
if setOfFrames is not None: if setOfFrames is not None:
self.listOfFrames = sorted(setOfFrames) self.listOfFrames = sorted(setOfFrames)
def __enter__(self):
self.fillBuffer()
return self
def getWH(self): def __exit__(self, type, value, traceback):
'''get width and height''' self.stop()
res, image = self.vc.read()
self.w = image.shape[1] def stop(self):
self.h = image.shape[0] self.thread.join()
return (self.w, self.h) self.vc.release()
def pop(self): def pop(self):
return self.buffer.get(block=True) return self.buffer.get(block=True)
@ -53,10 +53,6 @@ class VideoReader:
self.thread = threading.Thread(target=self.readFrames, args=()) self.thread = threading.Thread(target=self.readFrames, args=())
self.thread.start() self.thread.start()
def stop(self):
self.thread.join()
self.vc.release()
def readFrames(self): def readFrames(self):
'''Reads video from start to finish''' '''Reads video from start to finish'''
while self.lastFrame < self.endFrame: while self.lastFrame < self.endFrame:
@ -122,3 +118,11 @@ class VideoReader:
def getStartTime(self): def getStartTime(self):
return self.starttime return self.starttime
def getWH(self):
'''get width and height'''
if self.w is None or self.h is None:
res, image = self.vc.read()
self.w = image.shape[1]
self.h = image.shape[0]
return (self.w, self.h)

15
test.py Normal file
View File

@ -0,0 +1,15 @@
from Application.VideoReader import VideoReader
from Application.Config import Config
import os
fileName = "out.mp4"
dirName = os.path.join(os.path.dirname(__file__), "generate test footage")
config = {}
config["inputPath"] = os.path.join(dirName, fileName)
config["videoBufferLength"] = 100
with VideoReader(config) as reader:
while not reader.videoEnded():
framenumber, frame = reader.pop()
print(framenumber)