nodes/rhdunn-comfyui-audio-processing
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Allow the X and Y axis labels to be set.

Browse Source
This commit is contained in:
Reece H. Dunn
2024-08-06 20:17:11 +01:00
parent 771354e1ca
commit 2d0c29fbd9
1 changed files with 14 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
plot_nodes.py
+14 -6
View File
@@ -19,15 +19,17 @@ def plot2image():
class PlotWaveform:
@classmethod
def INPUT_TYPES(s):
return {"required": {"audio": ("AUDIO",),
"title": ("STRING", {"default": "Waveform"})}}
return {"required": {"audio": ("AUDIO",)},
"optional": {"title": ("STRING", {"default": "Waveform"}),
"xlabel": ("STRING", {"default": "Time (s)"}),
"ylabel": ("STRING", {"default": "Amplitude"}),}}
CATEGORY = "AudioProcessing/plot"
RETURN_TYPES = ("IMAGE", )
FUNCTION = "plot"
def plot(self, audio, title="Waveform"):
def plot(self, audio, title="Waveform", xlabel="Time (s)", ylabel="Amplitude"):
waveform = audio["waveform"].movedim(0, -1).numpy()
num_channels, num_frames, _ = waveform.shape
time_axis = torch.arange(0, num_frames) / audio["sample_rate"]
@@ -37,6 +39,8 @@ class PlotWaveform:
for c in range(num_channels):
axes[c].plot(time_axis, waveform[c], linewidth=1)
axes[c].grid(True)
axes[c].set_ylabel(ylabel)
axes[c].set_xlabel(xlabel)
if num_channels > 1:
axes[c].set_title(f"Channel {c+1}")
figure.suptitle(title)
@@ -46,15 +50,17 @@ class PlotWaveform:
class PlotSpectrogram:
@classmethod
def INPUT_TYPES(s):
return {"required": {"spectrogram": ("SPECT",),
"title": ("STRING", {"default": "Spectrogram"})}}
return {"required": {"spectrogram": ("SPECT",)},
"optional": {"title": ("STRING", {"default": "Spectrogram"}),
"xlabel": ("STRING", {"default": "Time (s)"}),
"ylabel": ("STRING", {"default": "Frequency (Hz)"}),}}
CATEGORY = "AudioProcessing/plot"
RETURN_TYPES = ("IMAGE", )
FUNCTION = "plot"
def plot(self, spectrogram, title="Spectrogram"):
def plot(self, spectrogram, title="Spectrogram", xlabel="Time (s)", ylabel="Frequency (Hz)"):
if len(spectrogram["spectrogram"].shape) == 2:
num_channels = 1
else:
@@ -68,6 +74,8 @@ class PlotSpectrogram:
else:
spect = spectrogram["spectrogram"][c].numpy()
axes[c].pcolormesh(spect, shading="gouraud")
axes[c].set_ylabel(ylabel)
axes[c].set_xlabel(xlabel)
if num_channels > 1:
axes[c].set_title(f"Channel {c+1}")
figure.suptitle(title)