separate script from class

background_heatmap.py

@@ -0,0 +1,52 @@
+import numpy as np
+import cv2
+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
+
+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.learning_rate = learning_rate
+        self.heatmap = np.array([])
+        self.backsub = cv2.createBackgroundSubtractorMOG2()
+        self.teach_background()
+        self.lastsum = self.to_floats(self.lastframe)
+
+    @staticmethod
+    def to_grayscale(frame):
+        return cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+
+    @staticmethod
+    def to_floats(frame):
+        return cv2.normalize(frame, None, alpha=0, beta=1, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F)
+
+    @staticmethod
+    def float_to_gray(frame):
+        return cv2.normalize(frame, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_8UC1)
+
+    @staticmethod
+    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)
+        self.heatmap = self.gray_to_heat(
+            self.float_to_gray(
+                self.to_floats(self.lastsum)
+            )
+        )
+
+    def teach_background(self):
+        for frame in self.frames:
+            if self.append_blur:
+                frame = cv2.blur(frame,(5,5))
+            self.lastframe = self.backsub.apply(frame, None, self.learning_rate)