Batch update after data-loss

utils/visual.py

@@ -1,4 +1,5 @@
 import numpy as np
+import librosa
 import matplotlib
 matplotlib.use('Agg')
 import matplotlib.pyplot as plt
@@ -14,6 +15,7 @@ def plot_alignment(alignment, info=None):
         xlabel += '\n\n' + info
     plt.xlabel(xlabel)
     plt.ylabel('Encoder timestep')
+    # plt.yticks(range(len(text)), list(text))
     plt.tight_layout()
     return fig
 
@@ -25,3 +27,28 @@ def plot_spectrogram(linear_output, audio):
     plt.colorbar()
     plt.tight_layout()
     return fig
+
+
+def visualize(alignment, spectrogram, stop_tokens, text, hop_length, CONFIG):
+    label_fontsize = 16
+    plt.figure(figsize=(16, 32))
+
+    plt.subplot(3, 1, 1)
+    plt.imshow(alignment.T, aspect="auto", origin="lower", interpolation=None)
+    plt.xlabel("Decoder timestamp", fontsize=label_fontsize)
+    plt.ylabel("Encoder timestamp", fontsize=label_fontsize)
+    plt.yticks(range(len(text)), list(text))
+    plt.colorbar()
+    
+    stop_tokens = stop_tokens.squeeze().detach().to('cpu').numpy()
+    plt.subplot(3, 1, 2)
+    plt.plot(range(len(stop_tokens)), list(stop_tokens))
+
+    plt.subplot(3, 1, 3)
+    librosa.display.specshow(spectrogram.T, sr=CONFIG.audio['sample_rate'],
+                             hop_length=hop_length, x_axis="time", y_axis="linear")
+    plt.xlabel("Time", fontsize=label_fontsize)
+    plt.ylabel("Hz", fontsize=label_fontsize)
+
+    plt.tight_layout()
+    plt.colorbar()