diff --git a/__init__.py b/__init__.py index 0be8d5c..963f924 100644 --- a/__init__.py +++ b/__init__.py @@ -1,11 +1,14 @@ # Copyright (C) 2024 Reece H. Dunn. SPDX-License-Identifier: GPL-3 from .io_nodes import LoadAudio +from .plot_nodes import PlotWaveform NODE_CLASS_MAPPINGS = { "ComfyAudio.LoadAudio": LoadAudio, + "ComfyAudio.PlotWaveform": PlotWaveform, } NODE_DISPLAY_NAME_MAPPINGS = { "ComfyAudio.LoadAudio": "Load Audio From Path", + "ComfyAudio.PlotWaveform": "Plot Waveform", } diff --git a/plot_nodes.py b/plot_nodes.py new file mode 100644 index 0000000..4f933a6 --- /dev/null +++ b/plot_nodes.py @@ -0,0 +1,43 @@ +# Copyright (C) 2024 Reece H. Dunn. SPDX-License-Identifier: GPL-3 + +import io +import matplotlib.pyplot as plt +import numpy as np +import torch + +from PIL import Image + + +def plot2image(): + buffer = io.BytesIO() + plt.savefig(buffer, format="png") + buffer.seek(0) + image = Image.open(buffer) + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) + + +class PlotWaveform: + @classmethod + def INPUT_TYPES(s): + return {"required": {"audio": ("AUDIO",), + "title": ("STRING", {"default": "Waveform"})}} + + CATEGORY = "AudioProcessing/plot" + + RETURN_TYPES = ("IMAGE", ) + FUNCTION = "plot" + + def plot(self, audio, title="Waveform"): + waveform = audio["waveform"].movedim(1, -1).numpy() + num_channels, num_frames, _ = waveform.shape + time_axis = torch.arange(0, num_frames) / audio["sample_rate"] + figure, axes = plt.subplots(num_channels, 1) + if num_channels == 1: + axes = [axes] + for c in range(num_channels): + axes[c].plot(time_axis, waveform[c], linewidth=1) + axes[c].grid(True) + if num_channels > 1: + axes[c].set_title(f"Channel {c+1}") + figure.suptitle(title) + return (plot2image(), )