97 lines
3.2 KiB
Python
97 lines
3.2 KiB
Python
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.max_history_len = 2
|
|
self.lastframes = []
|
|
self.learning_rate = learning_rate
|
|
self.heatmap = np.array([])
|
|
self.backsub = cv2.createBackgroundSubtractorMOG2()
|
|
self.teach_background()
|
|
self.reset_sum()
|
|
|
|
@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_heatmap(self):
|
|
self.heatmap = self.gray_to_heat(
|
|
self.float_to_gray(
|
|
self.to_floats(self.lastsum)
|
|
)
|
|
)
|
|
|
|
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))
|
|
# 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) |