Change plot function

requirements.txt

@@ -9,5 +9,4 @@ clu==0.0.7
 # pin Orbax to use Checkpointer
 orbax==0.0.2
 pytube
-pydub
-plotly==5.11.0
+pydub

utils.py

@@ -2,7 +2,6 @@
 import collections
 
 import pandas as pd
-import plotly.express as px
 import matplotlib.pyplot as plt
 from matplotlib.patches import Rectangle
 from PIL import Image
@@ -30,10 +29,9 @@ def dataframe_to_pianoroll_img(df):
     fig = plt.figure(figsize=(8, 5))
     ax = fig.add_subplot(111)
     ax.scatter(df.start_time, df.pitch, c="white")
-    n_colors = len(df.instrument.unique())
-    colorscale = px.colors.sample_colorscale("viridis", [n/(n_colors) for n in range(n_colors)])
-    print("colorscale", colorscale)
-    colordict = {f:colorscale[i] for i, f in enumerate(df.instrument.unique())}
+    colors = plt.get_cmap("tab20")
+    print(df["instrument"].unique())
+    colordict = {inst: colors[i] for i, inst in enumerate(df["instrument"].unique())}
     for _, row in df.iterrows():
         ax.add_patch(Rectangle((row["start_time"],
                                 row["pitch"]-0.4),