separate script from class
This commit is contained in:
52
background_heatmap.py
Normal file
52
background_heatmap.py
Normal file
@@ -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)
|
||||
@@ -1,96 +0,0 @@
|
||||
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, first_frame, append_blur=True, learning_rate=0.0003):
|
||||
self.append_blur = append_blur
|
||||
self.learning_rate = learning_rate
|
||||
self.heatmap = np.array([])
|
||||
self.backsub = cv2.createBackgroundSubtractorMOG2()
|
||||
# self.backsub = cv2.createBackgroundSubtractorKNN()
|
||||
# self.backsub = cv2.bgsegm_BackgroundSubtractorGSOC()
|
||||
# self.backsub = cv2.back()
|
||||
# # 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)
|
||||
if self.append_blur:
|
||||
frame = cv2.blur(first_frame,(5,5))
|
||||
self.lastframe = self.backsub.apply(frame, self.learning_rate)
|
||||
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)
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
projects_path = "C:\\UntrackedGit\\opencv_test\\him_projects"
|
||||
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]:
|
||||
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(2000)
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
# out.release()
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
37
heatmap_test.py
Normal file
37
heatmap_test.py
Normal file
@@ -0,0 +1,37 @@
|
||||
import cv2
|
||||
from video_loader import VideoLoader
|
||||
from background_heatmap import BackgroundHeatmap
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
||||
projects_path = "C:\\UntrackedGit\\opencv_test\\him_projects"
|
||||
loader = VideoLoader(projects_path)
|
||||
groups = loader.get_recordings_grouped()
|
||||
# out = cv2.VideoWriter('outpy.avi',cv2.VideoWriter_fourc\c('M','J','P','G'), 30, (512, 424), isColor=False)
|
||||
for group in groups:
|
||||
heatmaps = []
|
||||
for recording in group[1]:
|
||||
first = True
|
||||
frames = VideoLoader.extract_frames(recording)
|
||||
bgh = BackgroundHeatmap(VideoLoader.extract_frames(recording))
|
||||
for frame in frames:
|
||||
bgh.update(frame)
|
||||
cv2.imshow("IR", frame)
|
||||
# out.write(converted)
|
||||
# cv2.imshow("Heatmap", heatmap.heatmap)
|
||||
cv2.imshow("Background filter", bgh.lastframe)
|
||||
cv2.moveWindow("Background filter", 600, 100)
|
||||
if cv2.waitKey(1) & 0xFF == ord('q'):
|
||||
break
|
||||
heatmaps.append(bgh.heatmap)
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
for i, bgh in enumerate(heatmaps):
|
||||
imname = "Heatmap " + str(i)
|
||||
cv2.imshow(imname, bgh)
|
||||
cv2.moveWindow(imname, 500 * i, 0)
|
||||
cv2.waitKey(2000)
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
# out.release()
|
||||
cv2.destroyAllWindows()
|
||||
Reference in New Issue
Block a user