From b6c719e30d71fecb57b29f25f5fe73f2a7aa798d Mon Sep 17 00:00:00 2001
From: "Reece H. Dunn" <msclrhd@gmail.com>
Date: Mon, 19 Aug 2024 23:15:21 +0100
Subject: [PATCH] Plot Filter Bank

---
 README.md     |  9 +++++----
 __init__.py   |  4 +++-
 plot_nodes.py | 23 +++++++++++++++++++++++
 3 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 52bf8d3..c41bfb9 100644
--- a/README.md
+++ b/README.md
@@ -24,10 +24,11 @@
 
 ## Plotting
 
-| Name             | Description                                |
-|------------------|--------------------------------------------|
-| Plot Waveform    | Plots the waveform of an `AUDIO` object.   |
-| Plot Spectrogram | Plots the spectrogram of a `SPECT` object. |
+| Name                       | Description                                |
+|----------------------------|--------------------------------------------|
+| Plot Waveform              | Plots the waveform of an `AUDIO` object.   |
+| Plot Spectrogram           | Plots the spectrogram of a `SPECT` object. |
+| Plot Filter Bank           | Plots the `FILTER_BANK` object.            |
 
 ## Spectrogram
 
diff --git a/__init__.py b/__init__.py
index 8ca0a47..9a7da26 100644
--- a/__init__.py
+++ b/__init__.py
@@ -2,7 +2,7 @@
 
 from .fb_nodes import MelScaleFilterBank
 from .io_nodes import LoadAudio
-from .plot_nodes import (PlotWaveform, PlotSpectrogram)
+from .plot_nodes import (PlotWaveform, PlotSpectrogram, PlotFilterBank)
 from .spect_nodes import (
     Spectrogram, InverseSpectrogram,
     GriffinLim,
@@ -17,6 +17,7 @@ NODE_CLASS_MAPPINGS = {
     # Plotting
     "ComfyAudio.PlotWaveform": PlotWaveform,
     "ComfyAudio.PlotSpectrogram": PlotSpectrogram,
+    "ComfyAudio.PlotFilterBank": PlotFilterBank,
     # Spectrogram
     "ComfyAudio.Spectrogram": Spectrogram,
     "ComfyAudio.InverseSpectrogram": InverseSpectrogram,
@@ -32,6 +33,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     # Plotting
     "ComfyAudio.PlotWaveform": "Plot Waveform",
     "ComfyAudio.PlotSpectrogram": "Plot Spectrogram",
+    "ComfyAudio.PlotFilterBank": "Plot Filter Bank",
     # Spectrogram
     "ComfyAudio.Spectrogram": "Spectrogram",
     "ComfyAudio.InverseSpectrogram": "Inverse Spectrogram",
diff --git a/plot_nodes.py b/plot_nodes.py
index 791faa4..05a1cb8 100644
--- a/plot_nodes.py
+++ b/plot_nodes.py
@@ -87,3 +87,26 @@ class PlotSpectrogram:
         figure.suptitle(title)
         figure.tight_layout()
         return (plot2image(), )
+
+
+class PlotFilterBank:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"fbank": ("FILTER_BANK",)},
+                "optional": {"title": ("STRING", {"default": "Filter Bank"}),
+                             "xlabel": ("STRING", {"default": "filter bin"}),
+                             "ylabel": ("STRING", {"default": "frequency bin"}),}}
+
+    CATEGORY = "AudioProcessing/plot"
+
+    RETURN_TYPES = ("IMAGE", )
+    FUNCTION = "plot"
+
+    def plot(self, fbank, title="Filter Bank", xlabel="filter bin", ylabel="frequency bin"):
+        figure, axes = plt.subplots(1, 1)
+        axes.imshow(fbank['fbank'], origin="lower",aspect="auto", interpolation="nearest")
+        axes.set_ylabel(ylabel)
+        axes.set_xlabel(xlabel)
+        axes.set_title(title)
+        figure.tight_layout()
+        return (plot2image(), )