From e73fafcb9f9c0468058928ac123e2fe77d490797 Mon Sep 17 00:00:00 2001
From: Askill <ofjokg@gmail.com>
Date: Tue, 29 Dec 2020 10:41:55 +0100
Subject: [PATCH] added context manager for video reader

---
 Application/VideoReader.py | 36 ++++++++++++++++++++----------------
 test.py                    | 15 +++++++++++++++
 2 files changed, 35 insertions(+), 16 deletions(-)
 create mode 100644 test.py

diff --git a/Application/VideoReader.py b/Application/VideoReader.py
index 8ccfd9b..0963562 100644
--- a/Application/VideoReader.py
+++ b/Application/VideoReader.py
@@ -1,6 +1,3 @@
-from Application.Config import Config
-
-from datetime import datetime
 from queue import Queue
 
 import cv2
@@ -9,14 +6,13 @@ import os
 
 class VideoReader:
     listOfFrames = None
-    w = 0
-    h = 0
+    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:
@@ -30,13 +26,17 @@ class VideoReader:
         self.calcStartTime()
         if setOfFrames is not None:
             self.listOfFrames = sorted(setOfFrames)
+        
+    def __enter__(self):
+        self.fillBuffer()
+        return self
 
-    def getWH(self):
-        '''get width and height'''
-        res, image = self.vc.read()
-        self.w = image.shape[1]
-        self.h = image.shape[0]
-        return (self.w, self.h)
+    def __exit__(self, type, value, traceback):
+        self.stop()
+
+    def stop(self):
+        self.thread.join()
+        self.vc.release()
 
     def pop(self):
         return self.buffer.get(block=True)
@@ -53,10 +53,6 @@ class VideoReader:
             self.thread = threading.Thread(target=self.readFrames, args=())
         self.thread.start()
 
-    def stop(self):
-        self.thread.join()
-        self.vc.release()
-
     def readFrames(self):
         '''Reads video from start to finish'''
         while self.lastFrame < self.endFrame:
@@ -122,3 +118,11 @@ class VideoReader:
 
     def getStartTime(self):
         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)
diff --git a/test.py b/test.py
new file mode 100644
index 0000000..ee2500a
--- /dev/null
+++ b/test.py
@@ -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)