From 5775c6d7b4239e036f8506a7115b88793b6b6dd3 Mon Sep 17 00:00:00 2001
From: SayanoAI <28544690+SayanoAI@users.noreply.github.com>
Date: Mon, 16 Sep 2024 14:30:02 -0400
Subject: [PATCH] added diff map during training

---
 lib/train/utils.py | 4 ++--
 training_cli.py    | 1 +
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/lib/train/utils.py b/lib/train/utils.py
index 3cab854..c4fca31 100644
--- a/lib/train/utils.py
+++ b/lib/train/utils.py
@@ -187,7 +187,7 @@ def latest_checkpoint_path(dir_path, regex="G_*.pth"):
     return x
 
 
-def plot_spectrogram_to_numpy(spectrogram):
+def plot_spectrogram_to_numpy(spectrogram, cmap=None):
     global MATPLOTLIB_FLAG
     if not MATPLOTLIB_FLAG:
         import matplotlib
@@ -200,7 +200,7 @@ def plot_spectrogram_to_numpy(spectrogram):
     import numpy as np
 
     fig, ax = plt.subplots(figsize=(10, 2))
-    im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none")
+    im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none", cmap=cmap)
     plt.colorbar(im, ax=ax)
     plt.xlabel("Frames")
     plt.ylabel("Channels")
diff --git a/training_cli.py b/training_cli.py
index 4ae4676..e3f1d07 100644
--- a/training_cli.py
+++ b/training_cli.py
@@ -612,6 +612,7 @@ def train_and_evaluate(
                     "slice/mel_org": utils.plot_spectrogram_to_numpy(y_mel[0].data.cpu().numpy()),
                     "slice/mel_gen": utils.plot_spectrogram_to_numpy(y_hat_mel[0].data.cpu().numpy()),
                     "all/mel": utils.plot_spectrogram_to_numpy(mel[0].data.cpu().numpy()),
+                    "all/diff": utils.plot_spectrogram_to_numpy((y_mel[0]-y_mel[0]).abs().data.cpu().numpy(), cmap="hot")
                 }
 
                 with torch.no_grad():