added functionality of adding/removing plots of inlets

visualize_streams.py

@@ -40,14 +40,23 @@ class Inlet:
         self.name = info.name()
         self.channel_count = info.channel_count()
         self.uid = info.uid()
+        self.plt_ixs: List[int] = []
+        self.curves = None
+        self.enabled = False
 
-    def pull_and_plot(self, plot_time: float, plts: [pg.PlotItem]):
+    def set_enabled(self, enabled: bool):
+        self.enabled = enabled
+
+    def pull_and_plot(self, plot_time: float, gl: pg.GraphicsLayout, plts: [pg.PlotItem]):
         """Pull data from the inlet and add it to the plot.
         :param plot_time: lowest timestamp that's still visible in the plot
         :param plts: the plot the data should be shown on
         """
-        # We don't know what to do with a generic inlet, so we skip it.
-        pass
+        if not self.enabled:
+            return
+        else:
+            # We don't know what to do with a generic inlet, so we skip it.
+            pass
 
 
 class DataInlet(Inlet):
@@ -55,24 +64,29 @@ class DataInlet(Inlet):
     should be plotted as multiple lines."""
     dtypes = [[], np.float32, np.float64, None, np.int32, np.int16, np.int8, np.int64]
 
-    def __init__(self, info: pylsl.StreamInfo, gl: pg.GraphicsLayout, plts: [pg.PlotItem]):
+    def __init__(self, info: pylsl.StreamInfo):
         super().__init__(info)
         # calculate the size for our buffer, i.e. two times the displayed data
         bufsize = (2 * math.ceil(info.nominal_srate() * plot_duration), info.channel_count())
         self.buffer = np.empty(bufsize, dtype=self.dtypes[info.channel_format()])
-        empty = np.array([])
-        # create one curve object for each channel/line that will handle displaying the data
-        self.curves = [pg.PlotCurveItem(x=empty, y=empty, autoDownsample=True) for _ in range(self.channel_count)]
-        ch_ix = 0
-        for curve in self.curves:
-            gl.nextRow()
-            plts.append(gl.addPlot())
-            plts[-1].addItem(curve)
-            plts[-1].setXLink(plts[0])
-            plts[-1].setLabel('top', text=self.name + ' channel ' + str(ch_ix))
-            ch_ix = ch_ix + 1
-
-    def pull_and_plot(self, plot_time, plts):
+
+    def pull_and_plot(self, plot_time, gl: pg.GraphicsLayout, plts: [pg.PlotItem]):
+        if not self.enabled:
+            return
+        if len(self.plt_ixs) == 0:
+            empty = np.array([])
+            # create one curve object for each channel/line that will handle displaying the data
+            self.curves = [pg.PlotCurveItem(x=empty, y=empty, autoDownsample=True) for _ in range(self.channel_count)]
+            ch_ix = 0
+            for curve in self.curves:
+                gl.nextRow()
+                plts.append(gl.addPlot())
+                plts[-1].addItem(curve)
+                plts[-1].setXLink(plts[0])
+                plts[-1].setLabel('top', text=self.name + ' channel ' + str(ch_ix))
+                self.plt_ixs.append(len(plts) - 1)
+                ch_ix = ch_ix + 1
+
         # pull the data
         _, ts = self.inlet.pull_chunk(timeout=0.0,
                                       max_samples=self.buffer.shape[0],
@@ -110,7 +124,9 @@ class MarkerInlet(Inlet):
     def __init__(self, info: pylsl.StreamInfo):
         super().__init__(info)
 
-    def pull_and_plot(self, plot_time, plts):
+    def pull_and_plot(self, plot_time, gl: pg.GraphicsLayout, plts: [pg.PlotItem]):
+        if not self.enabled:
+            return
         # TODO: purge old markers
         strings, timestamps = self.inlet.pull_chunk(0)
         if timestamps:
@@ -151,25 +167,43 @@ def main():
     streams = pylsl.resolve_streams()
 
     def on_check_change(row, col, state):
-        # TODO actually update plots
         print("check changed: " + row + " " + str(state))
+        inlet_ix = 0
+        for inlet in inlets:
+            if inlet.name == row:
+                if state>0:
+                    inlet.set_enabled(True)
+                else:
+                    inlet.set_enabled(False)
+                    remove_plots_of_inlet(inlet)
+            inlet_ix += 1
+
+    def remove_plots_of_inlet(inlet):
+        for plt_ix in inlet.plt_ixs:
+            win.plots.removeItem(plts[plt_ix])
+        inlet.plt_ixs = []
 
     win.checkTable.sigStateChanged.connect(on_check_change)
 
     # iterate over found streams, creating specialized inlet objects that will
     # handle plotting the data
     for info in streams:
-        win.checkTable.addRow(info.name())
         if info.type() == 'Markers':
             if info.nominal_srate() != pylsl.IRREGULAR_RATE \
                     or info.channel_format() != pylsl.cf_string:
                 print('Invalid marker stream ' + info.name())
             print('Adding marker inlet: ' + info.name())
             inlets.append(MarkerInlet(info))
+            # give unique name to inlet and add to checkbox table
+            inlets[-1].name = str(len(inlets) - 1) + ' ' + inlets[-1].name
+            win.checkTable.addRow(inlets[-1].name)
         elif info.nominal_srate() != pylsl.IRREGULAR_RATE \
                 and info.channel_format() != pylsl.cf_string:
             print('Adding data inlet: ' + info.name())
-            inlets.append(DataInlet(info, win.plots, plts))
+            inlets.append(DataInlet(info))
+            # give unique name to inlet and add to checkbox table
+            inlets[-1].name = str(len(inlets) - 1) + ' ' + inlets[-1].name
+            win.checkTable.addRow(inlets[-1].name)
         else:
             print('Don\'t know what to do with stream ' + info.name())
 
@@ -188,7 +222,7 @@ def main():
         # Special handling of inlet types (markers, continuous data) is done in
         # the different inlet classes.
         for inlet in inlets:
-            inlet.pull_and_plot(mintime, plts)
+            inlet.pull_and_plot(mintime, win.plots, plts)
 
     # create a timer that will move the view every update_interval ms
     update_timer = QtCore.QTimer()