Making Heatmap class as dataclass. Making size, buffer and sigma as Heatmap attributes. Renaming init method into clear method.

1 files changed, 26 insertions, 31 deletions

diff --git a/src/argaze/AreaOfInterest/AOIFeatures.py b/src/argaze/AreaOfInterest/AOIFeatures.py
index cbe4483..8c684c0 100644
--- a/src/argaze/AreaOfInterest/AOIFeatures.py
+++ b/src/argaze/AreaOfInterest/AOIFeatures.py
@@ -8,6 +8,7 @@ __copyright__ = "Copyright 2023, Ecole Nationale de l'Aviation Civile (ENAC)"
 __license__ = "BSD"
 
 from typing import TypeVar, Tuple
+from dataclasses import dataclass, field
 import json
 
 from argaze import DataStructures
@@ -399,24 +400,34 @@ class TimeStampedAOIScenes(DataStructures.TimeStampedBuffer):
 HeatmapType = TypeVar('Heatmap', bound="Heatmap")
 # Type definition for type annotation convenience
 
+@dataclass
 class Heatmap():
     """Define image to draw heatmap."""
 
-    def __init__(self, size: tuple):
+    size: tuple = field(default=(1, 1))
+    """Size of heatmap image in pixels."""
 
-        self.__rX, self.__rY = size
+    buffer: int = field(default=0)
+    """Size of heatmap buffer (0 means no buffering)."""
+
+    sigma: float = field(default=0.05)
+    """Point spread factor."""
+
+    def __post_init__(self):
+
+        self.__rX, self.__rY = self.size
 
         # Init coordinates
         self.__Sx = numpy.linspace(0., self.__rX/self.__rY, self.__rX)
         self.__Sy = numpy.linspace(0., 1., self.__rY)
 
-        # Init heatmap
-        self.init()
+        # Init heatmap image
+        self.clear()
 
-    def point_spread(self, point: tuple, sigma: float):
+    def point_spread(self, point: tuple):
         """Draw gaussian point spread into image."""
 
-        div = -2 * sigma**2
+        div = -2 * self.sigma**2
 
         x = point[0] / self.__rY # we use rY not rX !!!
         y = point[1] / self.__rY
@@ -425,53 +436,37 @@ class Heatmap():
         dY2 = (self.__Sy - y)**2
 
         v_dX, v_dY = numpy.array(numpy.meshgrid(dX2, dY2)).reshape(2, -1)
-        
+
         return numpy.exp((v_dX + v_dY) / div).reshape(self.__rY, self.__rX)
 
-    def init(self, buffer_size: int = 0):
-        """Initialize heatmap image."""
+    def clear(self):
+        """Clear heatmap image."""
 
         self.__point_spread_sum = numpy.zeros((self.__rY, self.__rX))
         self.__point_spread_buffer = []
-        self.__point_spread_buffer_size = buffer_size
-
-    def update(self, point: tuple, sigma: float):
-        """Update heatmap image.
+        self.__point_spread_buffer_size = self.buffer
 
-        !!! danger
-                Call init() method before any update."""
+    def update(self, point: tuple):
+        """Update heatmap image."""
 
-        point_spread = self.point_spread(point, sigma)
+        point_spread = self.point_spread(point)
 
         # Sum point spread
         self.__point_spread_sum += point_spread
 
         # If point spread buffering enabled
-        if self.__point_spread_buffer_size > 0:
+        if self.buffer > 0:
 
             self.__point_spread_buffer.append(point_spread)
 
         # Remove oldest point spread buffer image
-        if len(self.__point_spread_buffer) > self.__point_spread_buffer_size:
+        if len(self.__point_spread_buffer) > self.buffer:
 
             self.__point_spread_sum -= self.__point_spread_buffer.pop(0)
 
         # Edit heatmap
         gray = (255 * self.__point_spread_sum / numpy.max(self.__point_spread_sum)).astype(numpy.uint8)
         self.__image = cv2.applyColorMap(gray, cv2.COLORMAP_JET)
-            
-    @property
-    def buffer(self) -> int:
-        """Get size of heatmap buffer."""
-
-        return self.__point_spread_buffer_size
-
-    @buffer.setter
-    def 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 image(self):