Adding a heatmap demo script in utils.

2 files changed, 64 insertions, 0 deletions

diff --git a/src/argaze/utils/README.md b/src/argaze/utils/README.md
index dc97c7a..0bd5012 100644
--- a/src/argaze/utils/README.md
+++ b/src/argaze/utils/README.md
@@ -51,4 +51,13 @@ Simulate gaze position using mouse pointer to illustrate gaze features.
 
 ```
 python ./src/argaze/utils/demo_gaze_features_run.py
+```
+
+# Heatmap demonstration
+
+Simulate gaze position using mouse pointer to produce heatmap picture.
+A picture is saved by pressing escape key.
+
+```
+python ./src/argaze/utils/demo_heatmap_run.py
 ```
\ No newline at end of file
diff --git a/src/argaze/utils/demo_heatmap_run.py b/src/argaze/utils/demo_heatmap_run.py
new file mode 100644
index 0000000..81aa98b
--- /dev/null
+++ b/src/argaze/utils/demo_heatmap_run.py
@@ -0,0 +1,55 @@
+#!/usr/bin/env python
+
+import os 
+
+from argaze.AreaOfInterest import AOIFeatures
+
+import numpy
+import cv2
+
+def main():
+
+	window_name = 'Heatmap'
+	frame_size = (800, 600)
+
+	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)
+
+	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)
+
+	# Attach mouse callback to window
+	cv2.setMouseCallback(window_name, on_mouse_event)
+
+	while True:
+
+		cv2.imshow(window_name, heatmap_matrix)
+
+		# 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)
+
+			break
+
+	# Stop frame display
+	cv2.destroyAllWindows()
+
+if __name__ == '__main__':
+
+    main()
\ No newline at end of file