Read videos from HIM project files

background_subtraction_heat.py

@@ -1,17 +1,25 @@
 import numpy as np
-from typing import List
 import cv2
+from typing import List
+from data.library import ArkiteData, ArkiteDetectionHEF
+from cv2util import to_8_bit_image
+import time
 
 class BackgroundHeatmap:
-    def __init__(self, capture):
+    def __init__(self, capture, bg_teach_iters=50):
         self.heatmap = np.array([])
         self.cap = capture
-        self.backsub = cv2.createBackgroundSubtractorMOG2()
+        # self.backsub = cv2.createBackgroundSubtractorMOG2()
+        self.backsub = cv2.createBackgroundSubtractorKNN()
         # self.backsub = cv2.bgsegm_BackgroundSubtractorGSOC()
         # self.backsub = cv2.back()
 
         # Fill up with first frame
         ret, frame = self.cap.read()
+        # Suppose the first frame is background, teach it for a few iterations
+        for i in range(bg_teach_iters):
+            ret, frame = self.cap.read()
+            self.backsub.apply(cv2.blur(frame,(5,5)))
         self.lastframe = self.backsub.apply(frame)
         self.lastsum = self.to_floats(self.lastframe)
 
@@ -33,7 +41,7 @@ class BackgroundHeatmap:
 
     def update(self):
         ret, frame = self.cap.read()
-        self.lastframe = self.backsub.apply(frame)
+        self.lastframe = self.backsub.apply(cv2.blur(frame,(5,5)))
         self.lastsum += self.to_floats(self.lastframe)
         self.heatmap = self.gray_to_heat(
             self.float_to_gray(
@@ -41,26 +49,43 @@ class BackgroundHeatmap:
             )
         )
 
-cap = cv2.VideoCapture(0)
-diffsum = BackgroundHeatmap(cap)
+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):
+            for recording in detection.recordings():
+                print("New recording: " + str(recording.name))
+                with recording.ir() as ir:
+                    for frame in ir.frame_sequence():
+                        converted = to_8_bit_image(frame, display_min=100, display_max=8000)
+                        # converted = cv2.applyColorMap(frame, colormap=cv2.COLORMAP_JET)
+                        cv2.imshow("IR", converted)
+                        if cv2.waitKey(1) & 0xFF == ord('q'):
+                            break
+                        # print("Showing frame...")
+                        # time.sleep(0.1)
 
-# Load diffsum up with first
-first = True
-while(True):
-    # Update heatmap
-    diffsum.update()
+    cv2.destroyAllWindows()
 
-    # Display the resulting frame
-    cv2.imshow('Heatmap', diffsum.heatmap)
-    cv2.imshow('Backsub', diffsum.lastframe)
+# # Load diffsum up with first
+# first = True
+# while(True):
+#     # Update heatmap
+#     diffsum.update()
 
-    if first:
-        cv2.moveWindow("Backsub", 1000, 100)
-        first = False
+#     # Display the resulting frame
+#     cv2.imshow('Heatmap', diffsum.heatmap)
+#     cv2.imshow('Backsub', diffsum.lastframe)
 
-    if cv2.waitKey(1) & 0xFF == ord('q'):
-        break
+#     if first:
+#         cv2.moveWindow("Backsub", 1000, 100)
+#         first = False
 
-# When everything done, release the capture
-diffsum.cap.release()
-cv2.destroyAllWindows()
+#     if cv2.waitKey(1) & 0xFF == ord('q'):
+#         break
+
+# # When everything done, release the capture
+# diffsum.cap.release()
+# cv2.destroyAllWindows()