Feat: support color settings

backend/bci_core/viz.py

@@ -120,20 +120,24 @@ def plot_confusion_matrix(y_true, y_pred):
     return disp.figure_
 
 
-def plot_states(time_range, pred_states, ax):
+def plot_states(time_range, pred_states, ax, colors=None):
     classes = np.unique(pred_states)
+    if colors is None:
+        colors = [plt.get_cmap('tab10')(i)[:3] for i in range(len(classes))]
     for i, c in enumerate(classes):
         ax.fill_between(np.linspace(*time_range, len(pred_states)), 0, 1,
-                        where=(pred_states == c), alpha=0.6, color=plt.get_cmap('tab10')(i)[:3])
+                        where=(pred_states == c), alpha=0.6, color=colors[i])
     return ax
 
 
-def plot_state_prob_with_cue(time_range, true_states, pred_probs, ax):
+def plot_state_prob_with_cue(time_range, true_states, pred_probs, ax, colors=None):
     # normalize
     ax.plot(np.linspace(*time_range, len(pred_probs)), pred_probs, color='k')
     # for each class, fill different colors
     classes = np.unique(true_states)
+    if colors is None:
+        colors = [plt.get_cmap('tab10')(i)[:3] for i in range(len(classes))]
     for i, c in enumerate(classes):
         ax.fill_between(np.linspace(*time_range, len(true_states)), 0, 1,
-                        where=(true_states == c), alpha=0.6, color=plt.get_cmap('tab10')(i)[:3])
+                        where=(true_states == c), alpha=0.6, color=colors[i])
     return ax