Merge pull request #3 from cvjena/develop

Develop - Add locations plotting as option to postprocess

README.md

@@ -74,7 +74,7 @@ Then use it in the `interpolate` function, and you are good to go.
 - [ ] Handle if not all values are given for a better schematic
 - [X] Add result images
 - [X] Add a function to draw triangulation onto the 2D canvas
-- [ ] Add a function to draw sensor locations onto the 2D canvas
+- [X] Add a function to draw sensor locations onto the 2D canvas
 - [X] Add the option to remove the area outside the canonical face model
 - [ ] Make a better interface for the channel names
 - [ ] Add function to create the according colorbar for matplotlib in the correct size

examples/data/fridlund.csv

@@ -1,7 +1,7 @@
-            ,DAO li ,OrbOr li ,Ment li ,Mass li ,Zyg li ,Llsup li ,OrbOc li ,lat Front li ,med Front li ,Corr li ,Deprsup li ,DAO re ,OrbOr re ,Ment re ,Mass re ,Zyg re ,Llsup re ,OrbOc re ,lat Front re ,med Front re ,Corr re ,Deprsup re
-angry       ,  9.87 ,   21.85 ,  16.80 ,   1.41 ,  1.10 ,    0.00 ,    6.89 ,        1.31 ,        0.93 ,   2.44 ,      2.74 ,  0.00 ,   39.11 ,  17.33 ,   2.04 ,  1.63 ,    0.00 ,    9.51 ,        0.71 ,        0.82 ,   2.27 ,      2.07
-suprised    ,  0.85 ,    0.79 ,   1.21 ,   0.46 ,  0.45 ,    0.61 ,    1.56 ,        4.71 ,        3.51 ,   2.26 ,      2.93 ,  0.73 ,    0.80 ,   1.12 ,   0.51 ,  0.52 ,    0.92 ,    3.01 ,        1.67 ,        3.95 ,   2.65 ,      2.46
-sad         ,  0.98 ,    1.23 ,  28.19 ,   0.63 ,  0.72 ,    3.49 ,    2.22 ,        3.09 ,        6.00 ,   7.91 ,      7.16 ,  1.00 ,    1.47 ,  28.97 ,   0.76 ,  0.80 ,    2.51 ,    3.41 ,        1.08 ,        6.72 ,   8.42 ,      7.04
-fearful     ,  8.68 ,    0.00 ,  11.20 ,   0.86 ,  0.92 ,    1.83 ,    1.58 ,        0.00 ,        5.65 ,   3.11 ,      2.99 ,  4.60 ,   11.33 ,   6.70 ,   1.11 ,  0.96 ,    2.05 ,    4.13 ,        1.60 ,        6.46 ,   3.68 ,      3.85
-happy       ,  3.14 ,   18.30 ,   8.37 ,   0.70 ,  0.00 ,    3.71 ,    6.05 ,        0.72 ,        0.65 ,   0.52 ,      0.33 ,  0.00 ,   15.04 ,  10.23 ,   0.92 ,  0.00 ,    3.31 ,    7.41 ,        0.60 ,        0.62 ,   0.58 ,      0.55
-distugusted ,  2.46 ,    8.31 ,   4.05 ,   0.77 ,  0.95 ,    1.79 ,    7.78 ,        3.65 ,        2.55 ,  10.57 ,      8.50 ,  1.90 ,    6.01 ,   3.92 ,   0.72 ,  0.81 ,    1.83 ,    8.97 ,        0.97 ,        2.39 ,  10.00 ,      8.13
+,DAO li,OrbOr li,Ment li,Mass li,Zyg li,Llsup li,OrbOc li,lat Front li,med Front li,Corr li,Deprsup li,DAO re,OrbOr re,Ment re,Mass re,Zyg re,Llsup re,OrbOc re,lat Front re,med Front re,Corr re,Deprsup re
+angry,9.87,21.85,16.80,1.41,1.10,0.00,6.89,1.31,0.93,2.44,2.74,0.00,39.11,17.33,2.04,1.63,0.00,9.51,0.71,0.82,2.27,2.07
+suprised,0.85,0.79,1.21,0.46,0.45,0.61,1.56,4.71,3.51,2.26,2.93,0.73,0.80,1.12,0.51,0.52,0.92,3.01,1.67,3.95,2.65,2.46
+sad,0.98,1.23,28.19,0.63,0.72,3.49,2.22,3.09,6.00,7.91,7.16,1.00,1.47,28.97,0.76,0.80,2.51,3.41,1.08,6.72,8.42,7.04
+fearful,8.68,0.00,11.20,0.86,0.92,1.83,1.58,0.00,5.65,3.11,2.99,4.60,11.33,6.70,1.11,0.96,2.05,4.13,1.60,6.46,3.68,3.85
+happy,3.14,18.30,8.37,0.70,0.00,3.71,6.05,0.72,0.65,0.52,0.33,0.00,15.04,10.23,0.92,0.00,3.31,7.41,0.60,0.62,0.58,0.55
+distugusted,2.46,8.31,4.05,0.77,0.95,1.79,7.78,3.65,2.55,10.57,8.50,1.90,6.01,3.92,0.72,0.81,1.83,8.97,0.97,2.39,10.00,8.13

src/electromyogram/__init__.py

@@ -6,12 +6,10 @@
     "Scheme",
     "colorize",
     "get_colormap",
-    "annotate_locations",
     "postprocess"
 ]
 
 from electromyogram.plot import (
-    annotate_locations,
     colorize,
     get_colormap,
     interpolate,

src/electromyogram/plot.py

@@ -1,4 +1,4 @@
-__all__ = ["interpolate", "plot_locations", "colorize", "get_colormap", "annotate_locations", "postprocess"]
+__all__ = ["interpolate", "plot_locations", "colorize", "get_colormap", "postprocess"]
 
 
 from dataclasses import dataclass
@@ -42,37 +42,6 @@ def get_outer(self, shape:tuple[int, int]) -> np.ndarray:
 
 face_model_data.close()
 
-def annotate_locations(
-    ax: plt.Axes,
-    scheme: Scheme,
-    dims: tuple[int, int] = (512, 512),
-    fontsize: int = 8,
-    marker_size: int = 3,
-    marker_color: str = "black",
-    draw_outer_hull: bool = False,
-    text_offset: tuple[int, int] = (0, 10),
-):
-    for emg_name, emg_loc in scheme.locations.items():
-        x, y = rel_to_abs(emg_loc[0], emg_loc[1], size=dims)
-        name = scheme.shortcuts.get(emg_name, emg_name)
-
-        ax.plot(x, y, marker=".", color=marker_color, markersize=marker_size)
-        ax.annotate(
-            name,
-            (x - text_offset[0], y - text_offset[1]),
-            xycoords="data",
-            va="bottom",
-            ha="center",
-            color="black",
-            fontsize=fontsize,
-        )
-
-    if draw_outer_hull:
-        for name, loc in scheme.outer_dict.items():
-            x, y = rel_to_abs(loc[0], loc[1], size=dims)
-            ax.plot(x, y, marker="D", color="green", markersize=marker_size // 2)
-
-
 def plot_locations(
     scheme: Scheme,
     shape: tuple[int, int] = (512, 512),
@@ -81,7 +50,8 @@ def plot_locations(
     radius: int = 7,
     color_circle: tuple[int, int, int] = (255, 105, 180),
     do_postprocess: bool = True,
-    draw_outerhull: bool = False, 
+    draw_outerhull: bool = False,
+    canvas: Optional[np.ndarray] = None,
 ) -> np.ndarray:
     """Plot the locations of the EMG values on a 2D canvas.
 
@@ -104,7 +74,8 @@ def plot_locations(
     np.ndarray
         The canvas with the plotted EMG values.
     """
-    canvas = np.full((shape[0], shape[1], 3), fill_value=255, dtype=np.uint8)
+    if canvas is None:
+        canvas = np.full((shape[0], shape[1], 3), fill_value=255, dtype=np.uint8)
     if do_postprocess:
         canvas = postprocess(canvas, remove_outer=True, draw_triangle=True, invert=True)
 
@@ -322,9 +293,13 @@ def postprocess(
     powermap: np.ndarray,
     remove_outer: bool = True,
     draw_triangle: bool = True,
+    draw_locations: bool = False,
     triangles_alpha: float = 0.3,
     invert: bool = False,
+    scheme: Optional[Scheme] = None,
 ) -> np.ndarray:
+    powermap = powermap.copy()
+
     # scale the points to the current shape
     points = (face_model.points * powermap.shape[0]).astype(np.int32)
     thickness = math.ceil(powermap.shape[0] / 512) # thickness of the lines optimized for a 512x512 canvas
@@ -334,14 +309,25 @@ def postprocess(
         lines = np.zeros_like(powermap)
         lines = cv2.polylines(lines, [points[tri] for tri in face_model.triangles], isClosed=True, color=color, thickness=thickness)
 
+        mask = np.zeros_like(powermap)
+        mask[lines != 0] = 1
+
         if invert:
             lines = cv2.bitwise_not(lines)
-
-        powermap = cv2.addWeighted(powermap, 1-triangles_alpha, lines, triangles_alpha, 0)
-
+
+        lines_masked = lines * mask
+        power_masked = powermap * mask
+
+        temp_blend = cv2.addWeighted(power_masked, 1-triangles_alpha, lines_masked, triangles_alpha, 0)
+        powermap[mask == 1] = temp_blend[mask == 1]
+
     if remove_outer:
         hull = cv2.convexHull(points, returnPoints=True)
         mask = cv2.drawContours(np.zeros(powermap.shape[:2]), [hull], 0, 1, -1)
         powermap[mask == 0] = 255 if powermap.ndim != 3 else [255, 255, 255]
 
+    if draw_locations:
+        assert scheme is not None, "scheme must not be None if draw_locations is True"
+        powermap = plot_locations(scheme=scheme, shape=powermap.shape[:2], do_postprocess=False, canvas=powermap)
+
     return powermap