added interactive data selector

teng_ml/data_preprocess.py

@@ -0,0 +1,127 @@
+import pandas as pd
+import numpy as np
+import scipy.signal as signal
+import matplotlib.pyplot as plt
+from time import sleep
+from random import choice as r_choice
+from sys import exit
+
+
+if __name__ == "__main__":
+    if __package__ is None:
+        # make relative imports work as described here: https://peps.python.org/pep-0366/#proposed-change
+        __package__ = "teng_ml"
+        import sys
+        from os import path
+        filepath = path.realpath(path.abspath(__file__))
+        sys.path.insert(0, path.dirname(path.dirname(filepath)))
+
+from .util.transform import Normalize
+from .util.data_loader import get_datafiles
+
+file = "/home/matth/Uni/TENG/teng_2/data/2023-06-28_foam_black_1_188mm_06V001.csv"
+
+class InteractiveDataSelector:
+    """
+    Helper class for "iterating" through selected peaks.
+    """
+    def __init__(self, out_name, out_dir, fig, ax):
+        self._out_dir = out_dir
+        self._out_name = out_name
+        self._fig = fig
+        self._ax = ax
+
+        self._fig.canvas.mpl_connect("button_press_event", lambda ev: self._fig_on_button_press(ev))
+        self._fig.canvas.mpl_connect("key_press_event", lambda ev: self._fig_on_key_press(ev))
+
+        self._splits_lines = None  # vlines
+        self._excludes_lines = None
+        self._excludes_areas = [] # list of areas
+
+        self._splits: list[int] = []
+        self._excludes: list[int] = []
+        self._mode = None  # split or exclude
+        self._set_mode("split")
+
+    def run(self):
+        while plt.fignum_exists(self._fig.number):
+            plt.pause(0.01)
+
+    def _fig_on_button_press(self, event):
+        if event.xdata in self._excludes or event.xdata in self._splits: return
+        if event.button == 1:  # left click, add position
+            if self._mode == "split":
+                self._splits.append(event.xdata)
+            else:
+                self._excludes.append(event.xdata)
+        elif event.button == 3:  # right click, undo
+            if self._mode == "split":
+                if len(self._splits) > 0:
+                    self._splits.pop()
+            else:
+                if len(self._excludes) > 0:
+                    self._excludes.pop()
+        self._update_lines()
+
+    def _fig_on_key_press(self, event):
+        if event.key == 'S':
+            self._set_mode("split")
+        elif event.key == 'e':
+            self._set_mode("exclude")
+
+    def _set_mode(self, mode):
+        help_str = "[(e)xclude - (S)plit]"
+        if mode == "split":
+            self._mode = "split"
+            fig.suptitle(f"-> split mode {help_str}")
+        else:
+            self._mode = "exclude"
+            fig.suptitle(f"-> exclude mode {help_str}")
+
+    def _update_lines(self):
+        print(self._splits, self._excludes)
+        ymin, ymax = self._ax.get_ylim()
+
+        if self._splits_lines is not None: self._splits_lines.remove()
+        self._splits_lines = self._ax.vlines(self._splits, ymin, ymax, color="b")
+
+        if self._excludes_lines is not None: self._excludes_lines.remove()
+        self._excludes_lines = self._ax.vlines(self._excludes, ymin, ymax, color="r")
+
+        for area in self._excludes_areas:
+            area.remove()
+        self._excludes_areas.clear()
+        excludes = self._excludes.copy()
+        if len(excludes) % 2 == 1: excludes.pop()  # only draw pairs
+        excludes.sort()
+        for i in range(1, len(excludes), 2):
+            self._excludes_areas.append(self._ax.axvspan(excludes[i-1], excludes[i], facecolor='r', alpha=0.3))
+
+        self._ax.set_ylim(ymin, ymax)  # reset, since margins are added to lines
+        self._fig.canvas.draw()
+
+    def _save_as_new_files(self):
+
+
+
+
+if __name__ == "__main__":
+    """
+    Peak identification:
+    plot, let user choose first, second, last and lowest peak for identification
+    """
+    df = pd.read_csv(file)
+    a = df.to_numpy()
+
+    # a2 = interpolate_to_linear_time()
+    # print(a2)
+    # exit()
+
+    vdata = Normalize(0, 1)(a[:,2])
+    plt.ion()
+    fig, ax = plt.subplots()
+    ax.plot(vdata)
+    ax.grid(True)
+    selector = InteractiveDataSelector("bla", "test", fig, ax)
+    selector.run()
+

teng_ml/peaks.py

@@ -10,7 +10,7 @@ from sys import exit
 if __name__ == "__main__":
     if __package__ is None:
         # make relative imports work as described here: https://peps.python.org/pep-0366/#proposed-change
-        __package__ = "teng-ml"
+        __package__ = "teng_ml"
         import sys
         from os import path
         filepath = path.realpath(path.abspath(__file__))