diff --git a/background_heatmap.py b/background_heatmap.py index 5095239..053221c 100644 --- a/background_heatmap.py +++ b/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): - 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) + def update_heatmap(self): self.heatmap = self.gray_to_heat( self.float_to_gray( self.to_floats(self.lastsum) ) ) - def teach_background(self): - for frame in self.frames: + def update(self, frame): + 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) \ No newline at end of file