From b8255e285d276a37351b296ce118797ec1716948 Mon Sep 17 00:00:00 2001 From: Bart Moyaers Date: Fri, 15 Nov 2019 10:59:02 +0100 Subject: [PATCH] group recordings by name --- background_subtraction_heat.py | 123 +++++++++++++++++++-------------- 1 file changed, 73 insertions(+), 50 deletions(-) diff --git a/background_subtraction_heat.py b/background_subtraction_heat.py index 10e370c..a512211 100644 --- a/background_subtraction_heat.py +++ b/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.createBackgroundSubtractorKNN() + 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) + # # Teach background filter 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.backsub.apply(frame) + self.lastframe = self.backsub.apply(cv2.blur(first_frame,(5,5))) # don't forget blur here if re applying blur 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): - ret, frame = self.cap.read() + def update(self, frame): 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): - 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) - cv2.destroyAllWindows() + grouper = RecordingGrouper(data) + groups = grouper.getRecordingGroups() -# # Load diffsum up with first -# first = True -# while(True): -# # Update heatmap -# diffsum.update() + 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) + 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() -# # Display the resulting frame -# 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() \ No newline at end of file + # cv2.destroyAllWindows()