provide ability to save traces and labels

setup.py

@@ -2,13 +2,10 @@
 
 setup(
     name="visplot",
-    install_requires=[
-        "vispy",
-        "numpy",
-    ],
+    install_requires=["pyqt5", "vispy", "numpy"],
     packages=find_packages(),
     py_modules=["visplot"],
     version=1.0,
     author="yhql",
     description="Side-channel traces visualizer",
-)
+)

visplot.py

@@ -1,23 +1,29 @@
-
 from itertools import cycle
 import numpy as np
+from PyQt5 import QtWidgets as qt
 from vispy import scene
 from vispy import color
-from vispy import util 
 
 
 class plot:
     """
     Fast plot of many large traces as a single object, using vispy. 
     """
-    MAX_HL = 12 
+
+    MAX_HL = 12
     BG_DARK = "#222"
     LBL_POS_DEFAULTX = 170
     LBL_POS_DEFAULTY = 40
-    LBL_SPACING = 16 
+    LBL_SPACING = 16
 
     def __init__(
-        self, curves=None, labels=None, bgcolor=BG_DARK, parent=None, dontrun=False
+        self,
+        curves=None,
+        labels=None,
+        bgcolor=BG_DARK,
+        parent=None,
+        dontrun=False,
+        no_controls=False,
     ):
         """ 
         :param icurves: input curve or list of curves
@@ -45,15 +51,16 @@ def __init__(
         self.x_axis.link_view(self.view)
         self.y_axis.link_view(self.view)
 
-        self.ctrl_pressed = False 
-        self.shift_pressed = False 
-        self.canvas.connect(self.on_key_press)
-        self.canvas.connect(self.on_key_release)
-        self.canvas.connect(self.on_mouse_press)
-        self.canvas.connect(self.on_mouse_release)
+        self.ctrl_pressed = False
+        self.shift_pressed = False
+
+        if not no_controls:
+            self.canvas.connect(self.on_key_press)
+            self.canvas.connect(self.on_key_release)
+            self.canvas.connect(self.on_mouse_press)
+            self.canvas.connect(self.on_mouse_release)
+            self.canvas.connect(self.on_mouse_move)
 
-        self.canvas.connect(self.on_mouse_move)
-
         if curves is not None:
             self.draw_curves(curves, labels)
 
@@ -64,7 +71,7 @@ def __init__(
     def clear(self):
         if self.line is not None:
             self.line.parent = None
-        self.selected_lines = [] 
+        self.selected_lines = []
         self.hl_labels = []
 
     def draw_curves(self, curves_, labels=None, clrmap="husl"):
@@ -98,26 +105,30 @@ def draw_curves(self, curves_, labels=None, clrmap="husl"):
         for i in range(size, nb_traces * size, size):
             connect[i - 1, 1] = i - 1
 
-        self.colors = np.ones((nb_traces*size,3),dtype=np.float32)
-        self.backup_colors = np.ones((nb_traces,3),dtype=np.float32)
+        self.colors = np.ones((nb_traces * size, 3), dtype=np.float32)
+        self.backup_colors = np.ones((nb_traces, 3), dtype=np.float32)
 
-        R_p = np.linspace(0.4,0.4,num=nb_traces)
-        G_p = np.linspace(0.5,0.3,num=nb_traces)
-        B_p = np.linspace(0.5,0.3,num=nb_traces)
+        R_p = np.linspace(0.4, 0.4, num=nb_traces)
+        G_p = np.linspace(0.5, 0.3, num=nb_traces)
+        B_p = np.linspace(0.5, 0.3, num=nb_traces)
 
-        self.colors[:,0] = np.repeat(R_p,size)
-        self.colors[:,1] = np.repeat(G_p,size)
-        self.colors[:,2] = np.repeat(B_p,size)
+        self.colors[:, 0] = np.repeat(R_p, size)
+        self.colors[:, 1] = np.repeat(G_p, size)
+        self.colors[:, 2] = np.repeat(B_p, size)
 
-        self.backup_colors[:,0] = R_p
-        self.backup_colors[:,1] = G_p
-        self.backup_colors[:,2] = B_p
+        self.backup_colors[:, 0] = R_p
+        self.backup_colors[:, 1] = G_p
+        self.backup_colors[:, 2] = B_p
 
-        self.line = scene.Line(pos=xy_curves, color=self.colors, parent=self.view.scene, connect=connect)
+        self.line = scene.Line(
+            pos=xy_curves, color=self.colors, parent=self.view.scene, connect=connect
+        )
 
-        self.selected_lines = [] 
+        self.selected_lines = []
         self.hl_labels = []
-        self.hl_colorset = cycle(color.get_colormap(clrmap)[np.linspace(0.0, 1.0, self.MAX_HL)])
+        self.hl_colorset = cycle(
+            color.get_colormap(clrmap)[np.linspace(0.0, 1.0, self.MAX_HL)]
+        )
 
         self.view.camera.set_range(x=(-1, size), y=(curves.min(), curves.max()))
 
@@ -129,88 +140,142 @@ def find_closest_line(self, x, y):
 
         tr = self.canvas.scene.node_transform(self.view.scene)
         # Canvas coordinates of clicked point
-        x1,y1,_,_ = tr.map((x,y))
+        x1, y1, _, _ = tr.map((x, y))
         # Canvas coordinates of upper right corner of bounding box
         # containing clicked point
-        x2,max_y,_,_ = tr.map((x+radius,y+radius))
+        x2, max_y, _, _ = tr.map((x + radius, y + radius))
 
-        _,min_y,_,_ = tr.map((x, y-radius))
+        _, min_y, _, _ = tr.map((x, y - radius))
         min_y, max_y = min(min_y, max_y), max(min_y, max_y)
 
         # Gather all segments left and right of the clicked point
         rx = int(round(x1))
-        tab = self.line.pos[:,rx-radius:rx+radius]
+        tab = self.line.pos[:, rx - radius : rx + radius]
 
-        # Find closest point, filtering out points whose 
+        # Find closest point, filtering out points whose
         # y-coordinate is outside the bounding box around
         # the clicked point
-        max_norm = 1000000 
-        imin = None 
-        f = np.array([x1,y1], dtype=np.float32)
-        for i,s in enumerate(tab):
+        max_norm = 1_000_000
+        imin = None
+        f = np.array([x1, y1], dtype=np.float32)
+        for i, s in enumerate(tab):
             for p in s:
-                if min_y<p[1]<max_y:
-                    t = np.linalg.norm(f-p)
+                if min_y < p[1] < max_y:
+                    t = np.linalg.norm(f - p)
                     if t < max_norm:
                         max_norm = t
                         imin = i
 
         # this is the index of the closest line
         # or None if there are no point in the
         # defined area
-        return imin 
+        return imin
 
-    def on_key_press(self, event):
-        if event.key == 'Control':
+    def _on_key_press(self, event):
+        if event.key == "Control":
             self.ctrl_pressed = True
-        if event.key == 'Shift':
+        if event.key == "Shift":
             self.shift_pressed = True
 
+    def on_key_press(self, event):
+        return self._on_key_press(event)
+
+    def _on_key_release(self, event):
+        if event.key == "Control":
+            self.ctrl_pressed = False
+        if event.key == "Shift":
+            self.shift_pressed = False
+
     def on_key_release(self, event):
-        if event.key == 'Control':
-            self.ctrl_pressed = False 
-        if event.key == 'Shift':
-            self.shift_pressed = False 
+        return self._on_key_release(event)
 
-    def on_mouse_press(self, event):
+    def _on_mouse_press(self, event):
         self.init_x, self.init_y = event.pos
 
-    def on_mouse_move(self,event):
+    def on_mouse_press(self, event):
+        return self._on_mouse_press(event)
+
+    def _on_mouse_move(self, event):
         if self.shift_pressed == True:
             if len(self.selected_lines) > 0:
                 # map to screen displacement
                 tr = self.canvas.scene.node_transform(self.view.scene)
-                x,_,_,_ = tr.map(event.pos)
-                init_x,_,_,_ = tr.map([self.init_x, self.init_y])
+                x, _, _, _ = tr.map(event.pos)
+                init_x, _, _, _ = tr.map([self.init_x, self.init_y])
                 delta_x = int(x - init_x)
                 for l in self.selected_lines:
-                    self.line.pos[l][:,1] = np.roll(self.line.pos[l][:,1],delta_x)
+                    self.line.pos[l][:, 1] = np.roll(self.line.pos[l][:, 1], delta_x)
                 self.init_x, self.init_y = event.pos
                 self.canvas.update()
 
-    def on_mouse_release(self, event):
-        ## ignore release when moving traces
+    def on_mouse_move(self, event):
+        return self._on_mouse_move(event)
+
+    def _on_mouse_release(self, event):
+        # ignore release when moving traces
         if self.shift_pressed:
             return
 
-        x,y = event.pos
+        x, y = event.pos
 
         # if released more than 3 pixels away from click (i.e. dragging), ignore
-        if not (abs(x-self.init_x)<3 and abs(y-self.init_y)<3):
+        if not (abs(x - self.init_x) < 3 and abs(y - self.init_y) < 3):
             return
 
-        closest_line = self.find_closest_line(x,y)
-        if closest_line is None:
-            return
+        if event.button == 1:
+            closest_line = self.find_closest_line(x, y)
+            if closest_line is None:
+                return
+
+            if self.ctrl_pressed:
+                self.multiple_select(closest_line)
+            else:
+                self.single_select(closest_line)
+        elif event.button == 2:
+            # save selected traces on right-click
+            self.save_traces_popup()
+
+    def save_traces_popup(self):
+        pt = qt.QLineEdit("", self.canvas.native)
+        pt.setPlaceholderText("Save traces as ...")
+        pt.show()
+        self.line_save = pt
+        # TODO: add a 'hasAcceptableInput'
+        pt.returnPressed.connect(self.save_traces)
+
+    def save_traces(self):
+        reshaped = self.line.pos[:, :, 1].reshape(self.shape_)
+        filename = self.line_save.text()
+        np.save(
+            filename + ".npy", reshaped 
+        )
+        print(f"Saved traces into {filename}.npy")
+        with open(filename + "_labels.txt", 'w') as out:
+            for label in self.labels:
+                print(label, file=out)
+
+        print(f"Saved labels into {filename}_labels.txt")
+        self.line_save.hide()
+
+    def on_mouse_release(self, event):
+        return self._on_mouse_release(event)
+
+    def get_trace_length(self):
+        if len(self.shape_) == 1:
+            return self.shape_[0]
+        return self.shape_[1]
 
-        if self.ctrl_pressed:
-            self.multiple_select(closest_line)
-        else:
-            self.single_select(closest_line)
-
     def _add_label(self, curve_index, new_color):
-        new_label = scene.Text(f"{self.labels[curve_index]}", color=new_color, anchor_x = 'left', parent=self.canvas.scene)
-        new_label.pos = self.LBL_POS_DEFAULTX, self.LBL_POS_DEFAULTY + self.LBL_SPACING * len(self.hl_labels)
+        new_label = scene.Text(
+            f"{self.labels[curve_index]}",
+            color=new_color,
+            anchor_x="left",
+            parent=self.canvas.scene,
+        )
+        new_label.pos = (
+            self.LBL_POS_DEFAULTX,
+            self.LBL_POS_DEFAULTY + self.LBL_SPACING * len(self.hl_labels),
+        )
         self.hl_labels.append((curve_index, new_label))
 
     def _del_label_from_curve_index(self, curve_index):
@@ -220,20 +285,23 @@ def _del_label_from_curve_index(self, curve_index):
 
         ## redraw text items
         for i, lbl in enumerate(self.hl_labels[idx:]):
-            lbl[1].pos = self.LBL_POS_DEFAULTX, self.LBL_POS_DEFAULTY + self.LBL_SPACING * (idx+i)
+            lbl[1].pos = (
+                self.LBL_POS_DEFAULTX,
+                self.LBL_POS_DEFAULTY + self.LBL_SPACING * (idx + i),
+            )
 
     def _find_label_from_curve_index(self, curve_index):
-        return list(map(lambda x:x[0], self.hl_labels)).index(curve_index)
+        return list(map(lambda x: x[0], self.hl_labels)).index(curve_index)
 
-    def _set_curve_color(self,n, new_color):
-        _, S = self.shape_
+    def _set_curve_color(self, n, new_color):
+        S = self.get_trace_length()
         a = n * S
-        self.colors[a:a+S] = np.repeat(new_color.rgb, S, axis=0)
+        self.colors[a : a + S] = np.repeat(new_color.rgb, S, axis=0)
 
-    def _restore_nth_curve_color(self,n):
-        _, S = self.shape_
+    def _restore_nth_curve_color(self, n):
+        S = self.get_trace_length()
         nnx = n * S
-        self.colors[nnx:nnx+S] = np.repeat([self.backup_colors[n]], S, axis=0)
+        self.colors[nnx : nnx + S] = np.repeat([self.backup_colors[n]], S, axis=0)
 
     def single_select(self, curve_index):
         # Unselect previously highlighted curves
@@ -247,12 +315,12 @@ def single_select(self, curve_index):
         self.hl_labels = []
 
         # Display its index/label
-        new_color = next(self.hl_colorset)               # Pick a new color
+        new_color = next(self.hl_colorset)  # Pick a new color
         self._add_label(curve_index, new_color)
-        self.selected_lines = [curve_index]              # Add this curve to the selected batch
-        self._set_curve_color(curve_index, new_color)    # Set its new color
+        self.selected_lines = [curve_index]  # Add this curve to the selected batch
+        self._set_curve_color(curve_index, new_color)  # Set its new color
 
-        self.line.set_data(color=self.colors)            # Update colors
+        self.line.set_data(color=self.colors)  # Update colors
 
     def multiple_select(self, curve_index):
         if curve_index in self.selected_lines:
@@ -264,17 +332,24 @@ def multiple_select(self, curve_index):
             self._restore_nth_curve_color(curve_index)
             self.selected_lines.remove(curve_index)
         else:
-            N,S = self.shape_
+            N, S = self.shape_
             new_color = next(self.hl_colorset)
             self._add_label(curve_index, new_color)
             self.selected_lines.append(curve_index)
             self._set_curve_color(curve_index, new_color)
-        
+
         self.line.set_data(color=self.colors)
 
+
 if __name__ == "__main__":
+    import sys
+
     N = 50
-    a = [i/10*np.sin(np.linspace(0.0+i/10,10.0+i/10,num=2000)) for i in range(N)]
+    a = [
+        i / 10 * np.sin(np.linspace(0.0 + i / 10, 10.0 + i / 10, num=2000))
+        for i in range(N)
+    ]
+
     v = plot(a, dontrun=True)
     v.multiple_select(4)
     v.multiple_select(7)