add separate video loader class

background_subtraction_heat.py

@@ -6,7 +6,6 @@ from video_loader import VideoLoader
 import numpy
 from data.library import ArkiteData
 from cv2util import to_8_bit_image
-from string import digits as DIGITS
 
 class BackgroundHeatmap:
     def __init__(self, first_frame, append_blur=True, learning_rate=0.0003):
@@ -53,73 +52,45 @@ class BackgroundHeatmap:
                 self.to_floats(self.lastsum)
             )
         )
+            
 
-class RecordingGrouper:
-    def __init__(self, data: ArkiteData):
-        self.data = data
-
-    def getRecordings(self):
-        for detection in data.detections(self.data.use_cases()):
-            for recording in detection.recordings():
-                yield recording
-
-    def getRecordingGroups(self):
-        groups = []
-        for recording in self.getRecordings():
-            # Check if last char is digit
-            name = recording.name
-            if name[-1] in DIGITS:
-                # Remove last two chars:
-                splitname = name[:-2]
-                group = next((x for x in groups if x[0] == splitname), None)
-                if group is not None:
-                    group[1].append(recording)
-                else:
-                    groups.append((splitname, [recording]))
-
-        return groups
-
-    
 if __name__ == '__main__':
 
     projects_path = "C:\\UntrackedGit\\opencv_test\\him_projects"
-    data = ArkiteData(projects_path, 1)
+    loader = VideoLoader(projects_path)
-
+    groups = loader.get_recordings_grouped()
-    grouper = RecordingGrouper(data)
+    # out = cv2.VideoWriter('outpy.avi',cv2.VideoWriter_fourcc('M','J','P','G'), 30, (512, 424), isColor=False)
-    groups = grouper.getRecordingGroups()
-
-    out = cv2.VideoWriter('outpy.avi',cv2.VideoWriter_fourcc('M','J','P','G'), 30, (512, 424), isColor=False)
-
     for group in groups:
         heatmaps = []
         for recording in group[1]:
-            print("New recording: " + str(recording.name))
+            first = True
-            with recording.ir() as ir:
+            for frame in VideoLoader.extract_frames(recording):
-                first = True
+                if first:
-                for frame in ir.frame_sequence():
+                    heatmap = BackgroundHeatmap(frame)
-                    if first:
+                    # Call an update on first frame again like HIM source
-                        heatmap = BackgroundHeatmap(frame)
+                    heatmap.update(frame)
-                        first = False
+                    first = False
-                    converted = to_8_bit_image(frame, display_min=100, display_max=8000)
+                heatmap.update(frame)
-                    heatmap.update(converted)
+                cv2.imshow("IR", frame)
-                    cv2.imshow("IR", converted)
+                # out.write(converted)
-                    # out.write(converted)
+                # cv2.imshow("Heatmap", heatmap.heatmap)
-                    cv2.imshow("Heatmap", heatmap.heatmap)
+                cv2.imshow("Background filter", heatmap.lastframe)
-                    cv2.moveWindow("Heatmap", 600, 100)
+                cv2.moveWindow("Background filter", 600, 100)
-                    if cv2.waitKey(1) & 0xFF == ord('q'):
+                if cv2.waitKey(1) & 0xFF == ord('q'):
-                        break
+                    break
-                    # print("Showing frame...")
+                # print("Showing frame...")
-                    # time.sleep(0.1)
+                # time.sleep(0.1)
-                heatmaps.append(heatmap.heatmap)
+            heatmaps.append(heatmap.heatmap)
-                cv2.destroyAllWindows()
+            cv2.destroyAllWindows()
+
         for i, heatmap in enumerate(heatmaps):
             imname = "Heatmap " + str(i)
             cv2.imshow(imname, heatmap)
             cv2.moveWindow(imname, 500 * i, 0)
-        cv2.waitKey()
+        cv2.waitKey(2000)
         cv2.destroyAllWindows()
 
     # out.release()
+    cv2.destroyAllWindows()
 
 
-    # cv2.destroyAllWindows()

recording_grouper.py

@@ -0,0 +1,27 @@
+from data.library import ArkiteData
+from string import digits as DIGITS
+
+class RecordingGrouper:
+    def __init__(self, data: ArkiteData):
+        self.data = data
+
+    def getRecordings(self):
+        for detection in self.data.detections(self.data.use_cases()):
+            for recording in detection.recordings():
+                yield recording
+
+    def getRecordingGroups(self):
+        groups = []
+        for recording in self.getRecordings():
+            # Check if last char is digit
+            name = recording.name
+            if name[-1] in DIGITS:
+                # Remove last two chars:
+                splitname = name[:-2]
+                group = next((x for x in groups if x[0] == splitname), None)
+                if group is not None:
+                    group[1].append(recording)
+                else:
+                    groups.append((splitname, [recording]))
+
+        return groups

video_loader.py

@@ -0,0 +1,52 @@
+from enum import Enum
+from data.library import ArkiteData, ArkiteDetection, ArkiteRecording
+from cv2util import to_8_bit_image
+from typing import Iterable
+import numpy
+from recording_grouper import RecordingGrouper
+
+class VideoType(Enum):
+    IR = 0
+    DEPTH = 1
+
+class VideoLoader:
+    def __init__(self, projects_path: str):
+        self.projects_path = projects_path
+        # Search folders with depth of 1:
+        self.data = ArkiteData(self.projects_path, 1)
+
+    def get_all_recordings(self) -> Iterable[ArkiteRecording]:
+        # Loop over projects
+        for detection in self.data.detections(self.data.use_cases()):
+            for recording in detection.recordings():
+                yield recording
+
+    def get_recordings_grouped(self):
+        grouper = RecordingGrouper(self.data)
+        return grouper.getRecordingGroups()
+    
+    @staticmethod
+    def extract_frames(recording: ArkiteRecording, video_type: VideoType = None) -> Iterable[numpy.ndarray]:
+        if video_type == VideoType.IR or video_type == None:
+            disp_min = 100
+            disp_max = 8000
+            with recording.ir() as ir:
+                frames = ir.frame_sequence()
+                for frame in frames:
+                    yield to_8_bit_image(frame, display_min=disp_min, display_max=disp_max)
+        elif video_type == VideoType.DEPTH:
+            # TODO: check if these values make sense for depth images.
+            disp_min = 100
+            disp_max = 8000
+            with recording.depth() as depth:
+                frames = depth.frame_sequence()
+                for frame in frames:
+                    yield to_8_bit_image(frame, display_min=disp_min, display_max=disp_max)
+        else:
+            raise NotImplementedError("VideoType \'{}\' not implemented".format(video_type))
+
+    # Yielding in sub method after "with" statement in top method does not work!
+    # @staticmethod
+    # def frames_to_8_bit(frame_sequence, disp_min=None, disp_max=None):
+    #     for frame in frame_sequence:
+    #         yield to_8_bit_image(frame, display_min=disp_min, display_max=disp_max)

videoloader_test.py

@@ -0,0 +1,12 @@
+from video_loader import VideoLoader
+import cv2
+from cv2util import to_8_bit_image
+
+projects_path = "C:\\UntrackedGit\\opencv_test\\him_projects"
+loader = VideoLoader(projects_path)
+
+for recording in loader.get_all_recordings():
+    frames = loader.extract_frames(recording)
+    for frame in frames:
+        cv2.imshow("test", frame)
+        cv2.waitKey(1)