From a58e04f445b5902215fce9f05988229214f4617f Mon Sep 17 00:00:00 2001 From: Théo de la Hogue Date: Wed, 24 May 2023 18:53:42 +0200 Subject: Adding heatmap feature inside AOIFrame class. --- src/argaze/AreaOfInterest/AOIFeatures.py | 74 ++++++++++++++++++++++++++---- src/argaze/utils/demo_gaze_features_run.py | 40 +++++----------- src/argaze/utils/demo_heatmap_run.py | 16 ++----- 3 files changed, 81 insertions(+), 49 deletions(-) (limited to 'src') diff --git a/src/argaze/AreaOfInterest/AOIFeatures.py b/src/argaze/AreaOfInterest/AOIFeatures.py index 8e3490c..edcfaf2 100644 --- a/src/argaze/AreaOfInterest/AOIFeatures.py +++ b/src/argaze/AreaOfInterest/AOIFeatures.py @@ -12,7 +12,7 @@ import json from argaze import DataStructures -import cv2 as cv +import cv2 import matplotlib.path as mpath import numpy from shapely.geometry import Polygon @@ -123,7 +123,7 @@ class AreaOfInterest(numpy.ndarray): Dst = numpy.array([[0., 0.], [1., 0.], [1., 1.], [0., 1.]]).astype(numpy.float32) - P = cv.getPerspectiveTransform(Src, Dst) + P = cv2.getPerspectiveTransform(Src, Dst) X = numpy.append(numpy.array(numpy.array(point) - Src_origin), [1.0]).astype(numpy.float32) Y = numpy.dot(P, X) @@ -146,7 +146,7 @@ class AreaOfInterest(numpy.ndarray): Dst_origin = Dst[0] Dst = (Dst - Dst_origin).reshape((len(Dst)), 2) - P = cv.getPerspectiveTransform(Src, Dst) + P = cv2.getPerspectiveTransform(Src, Dst) X = numpy.array([point[0], point[1], 1.0]).astype(numpy.float32) Y = numpy.dot(P, X) @@ -195,13 +195,13 @@ class AreaOfInterest(numpy.ndarray): # Draw form pixels = numpy.rint(self).astype(int) - cv.line(frame, pixels[-1], pixels[0], color, border_size) + cv2.line(frame, pixels[-1], pixels[0], color, border_size) for A, B in zip(pixels, pixels[1:]): - cv.line(frame, A, B, color, border_size) + cv2.line(frame, A, B, color, border_size) # Draw center center_pixel = numpy.rint(self.center).astype(int) - cv.circle(frame, center_pixel, 1, color, -1) + cv2.circle(frame, center_pixel, 1, color, -1) AOIFrameType = TypeVar('AOIFrame', bound="AOIFrame") # Type definition for type annotation convenience @@ -222,8 +222,8 @@ class AOIFrame(): self.__Sx = numpy.linspace(0., self.__rX/self.__rY, self.__rX) self.__Sy = numpy.linspace(0., 1., self.__rY) - # Init frame - self.__frame = numpy.zeros((self.__rY, self.__rX)) + # Init heatmap + self.heatmap_init() def point_spread(self, point: tuple, sigma: float): """Draw gaussian point spread into frame.""" @@ -240,6 +240,64 @@ class AOIFrame(): return numpy.exp((v_dX + v_dY) / div).reshape(self.__rY, self.__rX) + def heatmap_init(self, buffer_size: int = 0): + """Initialize heatmap matrix.""" + + self.__point_spread_sum = numpy.zeros((self.__rY, self.__rX)) + self.__point_spread_buffer = [] + self.__point_spread_buffer_size = buffer_size + + def heatmap_update(self, point: tuple, sigma: float): + """Update heatmap matrix. + + .. danger:: + Call heatmap_init() method before any update.""" + + point_spread = self.point_spread(point, sigma) + + # Sum point spread + self.__point_spread_sum += point_spread + + # If point spread buffering enabled + if self.__point_spread_buffer_size > 0: + + self.__point_spread_buffer.append(point_spread) + + # Remove oldest point spread buffer frame + if len(self.__point_spread_buffer) > self.__point_spread_buffer_size: + + self.__point_spread_sum -= self.__point_spread_buffer.pop(0) + + # Edit heatmap + heatmap_gray = (255 * self.__point_spread_sum / numpy.max(self.__point_spread_sum)).astype(numpy.uint8) + self.__heatmap_matrix = cv2.applyColorMap(heatmap_gray, cv2.COLORMAP_JET) + + + @property + def heatmap_buffer(self) -> int: + """Get size of heatmap buffer.""" + + return self.__point_spread_buffer_size + + @heatmap_buffer.setter + def heatmap_buffer(self, size: int): + """Set size of heatmap buffer (0 means no buffering).""" + + self.__point_spread_buffer = [] + self.__point_spread_buffer_size = size + + @property + def heatmap(self): + """Get heatmap matrix.""" + + try: + + return self.__heatmap_matrix + + except AttributeError: + + return numpy.zeros((self.__rY, self.__rX, 3)).astype(numpy.uint8) + AOISceneType = TypeVar('AOIScene', bound="AOIScene") # Type definition for type annotation convenience diff --git a/src/argaze/utils/demo_gaze_features_run.py b/src/argaze/utils/demo_gaze_features_run.py index 93ced82..cbaf720 100644 --- a/src/argaze/utils/demo_gaze_features_run.py +++ b/src/argaze/utils/demo_gaze_features_run.py @@ -59,13 +59,10 @@ def main(): gaze_position = GazeFeatures.GazePosition() screen_frame = AOIFeatures.AOIFrame(aoi_scene_projection['Screen'], window_size) - gaze_spread_sum = numpy.zeros((aoi_scene_image.shape[0], aoi_scene_image.shape[1])) - gaze_spread_buffer = [] - gaze_spread_buffer_size = 10 - heatmap_matrix = numpy.zeros(aoi_scene_image.shape, dtype=numpy.uint8) + screen_frame.heatmap_init() enable_heatmap = False - clear_sum_and_buffer = False + clear_heatmap = False enable_heatmap_buffer = False gaze_movement_identifier = { @@ -119,10 +116,7 @@ def main(): def on_mouse_event(event, x, y, flags, param): nonlocal gaze_position - nonlocal gaze_spread_sum - nonlocal gaze_spread_buffer - nonlocal heatmap_matrix - nonlocal clear_sum_and_buffer + nonlocal clear_heatmap nonlocal tm_probabilities nonlocal tm_density nonlocal raw_kc_analysis @@ -148,25 +142,14 @@ def main(): # Edit heatmap if enable_heatmap: - gaze_spread = screen_frame.point_spread(gaze_position.value, sigma=0.05) + # Clear heatmap + if clear_heatmap: - # Clear sum and buffer once - if clear_sum_and_buffer: - gaze_spread_sum = numpy.zeros((aoi_scene_image.shape[0], aoi_scene_image.shape[1])) - gaze_spread_buffer = [] - clear_sum_and_buffer = False + screen_frame.heatmap_init(10 if enable_heatmap_buffer else 0) + clear_heatmap = False - # Sum and and fill buffer - gaze_spread_sum += gaze_spread - gaze_spread_buffer.append(gaze_spread) - - # remove oldest gaze_spread buffer frame - if enable_heatmap_buffer and len(gaze_spread_buffer) > gaze_spread_buffer_size: - - gaze_spread_sum -= gaze_spread_buffer.pop(0) - - heatmap_gray = (255 * gaze_spread_sum / numpy.max(gaze_spread_sum)).astype(numpy.uint8) - heatmap_matrix = cv2.applyColorMap(heatmap_gray, cv2.COLORMAP_JET) + # Update heatmap + screen_frame.heatmap_update(gaze_position.value, sigma=0.05) else: @@ -271,7 +254,7 @@ def main(): # Draw gaze spread heatmap if enable_heatmap: - aoi_matrix = cv2.addWeighted(heatmap_matrix, 0.5, aoi_matrix, 1., 0) + aoi_matrix = cv2.addWeighted(screen_frame.heatmap, 0.5, aoi_matrix, 1., 0) else: @@ -448,8 +431,7 @@ def main(): enable_heatmap_buffer = not enable_heatmap_buffer - if enable_heatmap_buffer: - clear_sum_and_buffer = True + clear_heatmap = True # Enable Kc analysis with 'k' key if key_pressed == 107: diff --git a/src/argaze/utils/demo_heatmap_run.py b/src/argaze/utils/demo_heatmap_run.py index 81aa98b..3e6bb63 100644 --- a/src/argaze/utils/demo_heatmap_run.py +++ b/src/argaze/utils/demo_heatmap_run.py @@ -15,35 +15,27 @@ def main(): aoi = AOIFeatures.AreaOfInterest([[0, 0], [1, 0], [1, 1], [0, 1]]) aoi_frame = AOIFeatures.AOIFrame(aoi, frame_size) - gaze_spread_sum = numpy.zeros((frame_size[1], frame_size[0])) - heatmap_matrix = numpy.zeros((frame_size[1], frame_size[0]), dtype=numpy.uint8) + aoi_frame.heatmap_init() cv2.namedWindow(window_name, cv2.WINDOW_AUTOSIZE) # Update pointer position def on_mouse_event(event, x, y, flags, param): - nonlocal gaze_spread_sum - nonlocal heatmap_matrix - - gaze_spread = aoi_frame.point_spread((x, y), sigma=0.05) - gaze_spread_sum += gaze_spread - - heatmap_gray = (255 * gaze_spread_sum / numpy.max(gaze_spread_sum)).astype(numpy.uint8) - heatmap_matrix = cv2.applyColorMap(heatmap_gray, cv2.COLORMAP_JET) + aoi_frame.heatmap_update((x, y), sigma=0.05) # Attach mouse callback to window cv2.setMouseCallback(window_name, on_mouse_event) while True: - cv2.imshow(window_name, heatmap_matrix) + cv2.imshow(window_name, aoi_frame.heatmap) # Stop calibration by pressing 'Esc' key if cv2.waitKey(10) == 27: current_directory = os.path.dirname(os.path.abspath(__file__)) - cv2.imwrite(os.path.join(current_directory,'heatmap.png'), heatmap_matrix) + cv2.imwrite(os.path.join(current_directory,'heatmap.png'), aoi_frame.heatmap) break -- cgit v1.1