diff --git a/background_subtraction_heat.py b/background_subtraction_heat.py index cdc4a34..8417e9d 100644 --- a/background_subtraction_heat.py +++ b/background_subtraction_heat.py @@ -6,7 +6,6 @@ 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, first_frame, append_blur=True, learning_rate=0.0003): @@ -53,73 +52,45 @@ class BackgroundHeatmap: self.to_floats(self.lastsum) ) ) + -class RecordingGrouper: - def __init__(self, data: ArkiteData): - self.data = data - - def getRecordings(self): - for detection in data.detections(self.data.use_cases()): - for recording in detection.recordings(): - yield recording - - 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__': projects_path = "C:\\UntrackedGit\\opencv_test\\him_projects" - data = ArkiteData(projects_path, 1) - - grouper = RecordingGrouper(data) - groups = grouper.getRecordingGroups() - - out = cv2.VideoWriter('outpy.avi',cv2.VideoWriter_fourcc('M','J','P','G'), 30, (512, 424), isColor=False) - + loader = VideoLoader(projects_path) + groups = loader.get_recordings_grouped() + # out = cv2.VideoWriter('outpy.avi',cv2.VideoWriter_fourcc('M','J','P','G'), 30, (512, 424), isColor=False) 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) - # out.write(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() + first = True + for frame in VideoLoader.extract_frames(recording): + if first: + heatmap = BackgroundHeatmap(frame) + # Call an update on first frame again like HIM source + heatmap.update(frame) + first = False + heatmap.update(frame) + cv2.imshow("IR", frame) + # out.write(converted) + # cv2.imshow("Heatmap", heatmap.heatmap) + cv2.imshow("Background filter", heatmap.lastframe) + cv2.moveWindow("Background filter", 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.waitKey(2000) cv2.destroyAllWindows() # out.release() + cv2.destroyAllWindows() - # cv2.destroyAllWindows() diff --git a/recording_grouper.py b/recording_grouper.py new file mode 100644 index 0000000..bc8c0c3 --- /dev/null +++ b/recording_grouper.py @@ -0,0 +1,27 @@ +from data.library import ArkiteData +from string import digits as DIGITS + +class RecordingGrouper: + def __init__(self, data: ArkiteData): + self.data = data + + def getRecordings(self): + for detection in self.data.detections(self.data.use_cases()): + for recording in detection.recordings(): + yield recording + + 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 \ No newline at end of file diff --git a/video_loader.py b/video_loader.py new file mode 100644 index 0000000..f94341c --- /dev/null +++ b/video_loader.py @@ -0,0 +1,52 @@ +from enum import Enum +from data.library import ArkiteData, ArkiteDetection, ArkiteRecording +from cv2util import to_8_bit_image +from typing import Iterable +import numpy +from recording_grouper import RecordingGrouper + +class VideoType(Enum): + IR = 0 + DEPTH = 1 + +class VideoLoader: + def __init__(self, projects_path: str): + self.projects_path = projects_path + # Search folders with depth of 1: + self.data = ArkiteData(self.projects_path, 1) + + def get_all_recordings(self) -> Iterable[ArkiteRecording]: + # Loop over projects + for detection in self.data.detections(self.data.use_cases()): + for recording in detection.recordings(): + yield recording + + def get_recordings_grouped(self): + grouper = RecordingGrouper(self.data) + return grouper.getRecordingGroups() + + @staticmethod + def extract_frames(recording: ArkiteRecording, video_type: VideoType = None) -> Iterable[numpy.ndarray]: + if video_type == VideoType.IR or video_type == None: + disp_min = 100 + disp_max = 8000 + with recording.ir() as ir: + frames = ir.frame_sequence() + for frame in frames: + yield to_8_bit_image(frame, display_min=disp_min, display_max=disp_max) + elif video_type == VideoType.DEPTH: + # TODO: check if these values make sense for depth images. + disp_min = 100 + disp_max = 8000 + with recording.depth() as depth: + frames = depth.frame_sequence() + for frame in frames: + yield to_8_bit_image(frame, display_min=disp_min, display_max=disp_max) + else: + raise NotImplementedError("VideoType \'{}\' not implemented".format(video_type)) + + # Yielding in sub method after "with" statement in top method does not work! + # @staticmethod + # def frames_to_8_bit(frame_sequence, disp_min=None, disp_max=None): + # for frame in frame_sequence: + # yield to_8_bit_image(frame, display_min=disp_min, display_max=disp_max) \ No newline at end of file diff --git a/videoloader_test.py b/videoloader_test.py new file mode 100644 index 0000000..f135adc --- /dev/null +++ b/videoloader_test.py @@ -0,0 +1,12 @@ +from video_loader import VideoLoader +import cv2 +from cv2util import to_8_bit_image + +projects_path = "C:\\UntrackedGit\\opencv_test\\him_projects" +loader = VideoLoader(projects_path) + +for recording in loader.get_all_recordings(): + frames = loader.extract_frames(recording) + for frame in frames: + cv2.imshow("test", frame) + cv2.waitKey(1) \ No newline at end of file