group recordings by name

background_subtraction_heat.py

@@ -1,26 +1,26 @@
 import numpy as np
 import cv2
-from typing import List
-from data.library import ArkiteData, ArkiteDetectionHEF
-from cv2util import to_8_bit_image
 import time
+from typing import List, Iterable
+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, capture, bg_teach_iters=50):
+    def __init__(self, first_frame):
         self.heatmap = np.array([])
-        self.cap = capture
+        self.backsub = cv2.createBackgroundSubtractorMOG2()
-        # self.backsub = cv2.createBackgroundSubtractorMOG2()
+        # self.backsub = cv2.createBackgroundSubtractorKNN()
-        self.backsub = cv2.createBackgroundSubtractorKNN()
         # self.backsub = cv2.bgsegm_BackgroundSubtractorGSOC()
         # self.backsub = cv2.back()
-
+        # # Teach background filter for a few iterations
-        # Fill up with first frame
+        # for i in range(bg_teach_iters):
-        ret, frame = self.cap.read()
+        #     ret, frame = self.cap.read()
-        # Suppose the first frame is background, teach it for a few iterations
+        #     # self.backsub.apply(cv2.blur(frame,(5,5)))
-        for i in range(bg_teach_iters):
+        #     self.backsub.apply(frame)
-            ret, frame = self.cap.read()
+        self.lastframe = self.backsub.apply(cv2.blur(first_frame,(5,5))) # don't forget blur here if re applying blur
-            self.backsub.apply(cv2.blur(frame,(5,5)))
-        self.lastframe = self.backsub.apply(frame)
         self.lastsum = self.to_floats(self.lastframe)
 
     @staticmethod
@@ -39,8 +39,7 @@ class BackgroundHeatmap:
     def gray_to_heat(frame):
         return cv2.applyColorMap(frame, cv2.COLORMAP_JET)
 
-    def update(self):
+    def update(self, frame):
-        ret, frame = self.cap.read()
         self.lastframe = self.backsub.apply(cv2.blur(frame,(5,5)))
         self.lastsum += self.to_floats(self.lastframe)
         self.heatmap = self.gray_to_heat(
@@ -49,43 +48,67 @@ class BackgroundHeatmap:
             )
         )
 
+class RecordingGrouper:
+    def __init__(self, data: ArkiteData):
+        self.data = data
+
+    def getRecordings(self):
+        for uc in self.data.use_cases():
+            for detection in data.detections(uc):
+                return detection.recordings()
+
+    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__':
-    # cap = cv2.VideoCapture(0)
+
     projects_path = "C:\\UntrackedGit\\opencv_test\\him_projects"
     data = ArkiteData(projects_path, 1)
-    for uc in data.use_cases():
+
-        for detection in data.detections(uc):
+    grouper = RecordingGrouper(data)
-            for recording in detection.recordings():
+    groups = grouper.getRecordingGroups()
+
+    for group in groups:
+        heatmaps = []
+        for recording in group[1]:
             print("New recording: " + str(recording.name))
             with recording.ir() as ir:
+                first = True
                 for frame in ir.frame_sequence():
+                    if first:
+                        heatmap = BackgroundHeatmap(frame)
+                        first = False
                     converted = to_8_bit_image(frame, display_min=100, display_max=8000)
-                        # converted = cv2.applyColorMap(frame, colormap=cv2.COLORMAP_JET)
+                    heatmap.update(converted)
                     cv2.imshow("IR", converted)
+                    cv2.imshow("Heatmap", heatmap.heatmap)
+                    cv2.moveWindow("Heatmap", 600, 100)
                     if cv2.waitKey(1) & 0xFF == ord('q'):
                         break
                     # print("Showing frame...")
                     # time.sleep(0.1)
-
+                heatmaps.append(heatmap.heatmap)
+                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.destroyAllWindows()
 
-# # Load diffsum up with first
-# first = True
-# while(True):
-#     # Update heatmap
-#     diffsum.update()
 
-#     # Display the resulting frame
+    # cv2.destroyAllWindows()
-#     cv2.imshow('Heatmap', diffsum.heatmap)
-#     cv2.imshow('Backsub', diffsum.lastframe)
-
-#     if first:
-#         cv2.moveWindow("Backsub", 1000, 100)
-#         first = False
-
-#     if cv2.waitKey(1) & 0xFF == ord('q'):
-#         break
-
-# # When everything done, release the capture
-# diffsum.cap.release()
-# cv2.destroyAllWindows()