add group backfroundfilter

background_heatmap.py

@@ -11,12 +11,13 @@ class BackgroundHeatmap:
     def __init__(self, frames, append_blur=True, learning_rate=0.0003):
         self.append_blur = append_blur
         self.frames = frames
-        self.lastframe = None
+        self.max_history_len = 2
+        self.lastframes = []
         self.learning_rate = learning_rate
         self.heatmap = np.array([])
         self.backsub = cv2.createBackgroundSubtractorMOG2()
         self.teach_background()
-        self.lastsum = self.to_floats(self.lastframe)
+        self.reset_sum()
 
     @staticmethod
     def to_grayscale(frame):
@@ -34,19 +35,63 @@ class BackgroundHeatmap:
     def gray_to_heat(frame):
         return cv2.applyColorMap(frame, cv2.COLORMAP_JET)
 
-    def update(self, frame):
+    def update_heatmap(self):
-        if self.append_blur:
-            frame = cv2.blur(frame,(5,5))
-        self.backsub.apply(frame, self.lastframe, self.learning_rate)
-        self.lastsum += self.to_floats(self.lastframe)
         self.heatmap = self.gray_to_heat(
             self.float_to_gray(
                 self.to_floats(self.lastsum)
             )
         )
 
-    def teach_background(self):
+    def update(self, frame):
-        for frame in self.frames:
+        if self.append_blur:
+            frame = cv2.blur(frame,(5,5))
+        self.add_frame_history(self.backsub.apply(frame, None, self.learning_rate))
+        self.lastsum += self.to_floats(self.lastframe)
+        self.update_heatmap()
+
+    def teach_background(self, frames=None):
+        if frames is None:
+            frames = self.frames
+        for frame in frames:
             if self.append_blur:
                 frame = cv2.blur(frame,(5,5))
-            self.lastframe = self.backsub.apply(frame, None, self.learning_rate)
+            # TODO: is it logical here to use the same learning rate?
+            self.add_frame_history(self.backsub.apply(frame, None, self.learning_rate))
+
+    def add_frame_history(self, frame):
+        if len(self.lastframes) == self.max_history_len:
+            self.lastframes.pop(0)
+        self.lastframes.append(frame)
+
+    def reset_sum(self):
+        self.lastsum = self.to_floats(self.lastframe)
+
+    def reset(self):
+        self.reset_sum()
+        self.update_heatmap()
+
+    @property
+    def lastframe(self):
+        return self.lastframes[-1]
+
+    @property
+    def bgf_diff(self):
+        # make binary map
+        first = cv2.threshold(self.lastframes[1], 60, 255, cv2.THRESH_BINARY)[1]
+        second = cv2.threshold(self.lastframes[0], 60, 255, cv2.THRESH_BINARY)[1]
+        return cv2.subtract(first, second)
+
+class BackgroundHeatmapFromGroup(BackgroundHeatmap):
+    def __init__(self, group, append_blur=True, learning_rate=0.0003):
+        self.group = group
+        assert len(self.group[1]) is not 0
+        # Init with first recording
+        frames = VideoLoader.extract_frames(self.group[1][0])
+        super().__init__(frames, append_blur=append_blur, learning_rate=learning_rate)
+        self.teachRestOfVideosInGroup()
+    
+    def teachRestOfVideosInGroup(self):
+        # Slice away first recording since it was passed to super constructor already
+        for recording in self.group[1][1:]:
+            frames = VideoLoader.extract_frames(recording)
+            self.teach_background(frames=frames)