285 lines
8.2 KiB
Python
285 lines
8.2 KiB
Python
import io
|
|
|
|
import librosa
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import torch
|
|
import torchaudio.functional as F
|
|
from PIL import Image
|
|
|
|
from .node_def import (
|
|
BASE_NODE_CATEGORY,
|
|
MAX_SAFE_INT,
|
|
MAX_SAMPLE_RATE,
|
|
AudioData,
|
|
SpectrogramData,
|
|
)
|
|
|
|
NODE_CATEGORY = BASE_NODE_CATEGORY + "/visualize"
|
|
|
|
|
|
# PIL to Tensor
|
|
# refer: https://github.com/a1lazydog/ComfyUI-AudioScheduler/blob/main/nodes.py
|
|
def pil2tensor(images):
|
|
tensors = [
|
|
torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
|
for image in images
|
|
]
|
|
return torch.concat(tensors, dim=0)
|
|
|
|
|
|
class PlotWaveForm:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"audio": ("AUDIO",),
|
|
"title": ("STRING", {"default": "Waveform"}),
|
|
},
|
|
}
|
|
|
|
CATEGORY = NODE_CATEGORY
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("graph_image",)
|
|
FUNCTION = "plot_waveform"
|
|
|
|
def plot_waveform(self, audio: AudioData, title: str):
|
|
# refer: https://pytorch.org/tutorials/beginner/audio_preprocessing_tutorial.html
|
|
waveform = audio["waveform"].numpy()
|
|
|
|
num_batch, num_channels, num_frames = waveform.shape
|
|
images = []
|
|
for b in range(num_batch):
|
|
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[b][c], linewidth=1)
|
|
axes[c].grid(True)
|
|
if num_channels > 1:
|
|
axes[c].set_ylabel(f"Channel {c + 1}")
|
|
figure.suptitle(title)
|
|
|
|
# Create an in-memory buffer to store the image
|
|
buffer = io.BytesIO()
|
|
|
|
# Save the plot to the in-memory buffer as a PNG
|
|
plt.savefig(buffer, format="png")
|
|
buffer.seek(0)
|
|
|
|
# Create a Pillow Image object
|
|
image = Image.open(buffer)
|
|
images.append(image)
|
|
|
|
return (pil2tensor(images),)
|
|
|
|
|
|
class PlotSpecgram:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"audio": ("AUDIO",),
|
|
"title": ("STRING", {"default": "Spectrogram"}),
|
|
},
|
|
}
|
|
|
|
CATEGORY = NODE_CATEGORY
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("graph_image",)
|
|
FUNCTION = "plot_specgram"
|
|
|
|
def plot_specgram(self, audio: AudioData, title: str):
|
|
# refer: https://pytorch.org/tutorials/beginner/audio_preprocessing_tutorial.html
|
|
waveform = audio["waveform"].numpy()
|
|
|
|
num_batch, num_channels, _ = waveform.shape
|
|
# not used in the tutorial
|
|
# time_axis = torch.arange(0, num_frames) / audio["sample_rate"]
|
|
images = []
|
|
for b in range(num_batch):
|
|
figure, axes = plt.subplots(num_channels, 1)
|
|
if num_channels == 1:
|
|
axes = [axes]
|
|
for c in range(num_channels):
|
|
axes[c].specgram(waveform[b][c], Fs=audio["sample_rate"])
|
|
figure.suptitle(title)
|
|
|
|
# Create an in-memory buffer to store the image
|
|
buffer = io.BytesIO()
|
|
|
|
# Save the plot to the in-memory buffer as a PNG
|
|
plt.savefig(buffer, format="png")
|
|
buffer.seek(0)
|
|
|
|
# Create a Pillow Image object
|
|
image = Image.open(buffer)
|
|
images.append(image)
|
|
|
|
return (pil2tensor(images),)
|
|
|
|
|
|
class PlotSpectrogram:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"spec": ("SPEC",),
|
|
"title": ("STRING", {"default": "Spectrogram (db)"}),
|
|
"ylabel": ("STRING", {"default": "freq_bin"}),
|
|
},
|
|
}
|
|
|
|
CATEGORY = NODE_CATEGORY
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("graph_image",)
|
|
FUNCTION = "plot_spectrogram"
|
|
|
|
def plot_spectrogram(self, spec: SpectrogramData, title: str, ylabel: str):
|
|
# refer: https://pytorch.org/tutorials/beginner/audio_preprocessing_tutorial.html
|
|
|
|
num_batch, _, _ = spec.waveform.shape
|
|
|
|
images = []
|
|
for b in range(num_batch):
|
|
fig, axs = plt.subplots(1, 1)
|
|
axs.set_title(title)
|
|
axs.set_ylabel(ylabel)
|
|
axs.set_xlabel("frame")
|
|
im = axs.imshow(
|
|
librosa.power_to_db(spec.waveform[b][0].numpy()),
|
|
origin="lower",
|
|
aspect="auto",
|
|
)
|
|
fig.colorbar(im, ax=axs)
|
|
|
|
# Create an in-memory buffer to store the image
|
|
buffer = io.BytesIO()
|
|
|
|
# Save the plot to the in-memory buffer as a PNG
|
|
plt.savefig(buffer, format="png")
|
|
buffer.seek(0)
|
|
|
|
# Create a Pillow Image object
|
|
image = Image.open(buffer)
|
|
images.append(image)
|
|
|
|
return (pil2tensor(images),)
|
|
|
|
|
|
class PlotMelFilterBank:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"n_fft": ("INT", {"default": 256, "min": 0, "max": MAX_SAFE_INT}),
|
|
"n_mels": ("INT", {"default": 64, "min": 0, "max": MAX_SAFE_INT}),
|
|
"sample_rate": ("INT", {"default": 32000, "min": 0, "max": MAX_SAMPLE_RATE}),
|
|
"title": ("STRING", {"default": "Filter bank"}),
|
|
},
|
|
}
|
|
|
|
CATEGORY = NODE_CATEGORY
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("graph_image",)
|
|
FUNCTION = "plot_mel_fbank"
|
|
|
|
def plot_mel_fbank(self, n_fft, n_mels, sample_rate, title: str):
|
|
fbank = F.melscale_fbanks(
|
|
int(n_fft // 2 + 1),
|
|
n_mels=n_mels,
|
|
f_min=0.0,
|
|
f_max=sample_rate / 2.0,
|
|
sample_rate=sample_rate,
|
|
norm="slaney",
|
|
)
|
|
fig, axs = plt.subplots(1, 1)
|
|
axs.set_title(title)
|
|
axs.imshow(fbank, aspect="auto")
|
|
axs.set_ylabel("frequency bin")
|
|
axs.set_xlabel("mel bin")
|
|
# Create an in-memory buffer to store the image
|
|
buffer = io.BytesIO()
|
|
|
|
# Save the plot to the in-memory buffer as a PNG
|
|
plt.savefig(buffer, format="png")
|
|
buffer.seek(0)
|
|
|
|
# Create a Pillow Image object
|
|
image = Image.open(buffer)
|
|
|
|
return (pil2tensor([image]),)
|
|
|
|
|
|
class PlotPitch:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"audio": ("AUDIO",)}}
|
|
|
|
CATEGORY = NODE_CATEGORY
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("graph_image",)
|
|
FUNCTION = "plot_pitch"
|
|
|
|
def plot_pitch(self, audio: AudioData):
|
|
pitch = F.detect_pitch_frequency(audio["waveform"], audio["sample_rate"])
|
|
|
|
num_batch, _, _ = audio["waveform"].shape
|
|
images = []
|
|
|
|
for b in range(num_batch):
|
|
figure, axis = plt.subplots(1, 1)
|
|
axis.set_title("Pitch Feature")
|
|
axis.grid(True)
|
|
|
|
end_time = audio["waveform"].shape[-1] / audio["sample_rate"]
|
|
time_axis = torch.linspace(0, end_time, audio["waveform"].shape[-1])
|
|
axis.plot(
|
|
time_axis, audio["waveform"][b][0], linewidth=1, color="gray", alpha=0.3
|
|
)
|
|
|
|
axis2 = axis.twinx()
|
|
time_axis = torch.linspace(0, end_time, pitch.shape[-1])
|
|
ln2 = axis2.plot(
|
|
time_axis, pitch[b][0], linewidth=2, label="Pitch", color="green"
|
|
)
|
|
|
|
axis2.legend(loc=0)
|
|
|
|
# Create an in-memory buffer to store the image
|
|
buffer = io.BytesIO()
|
|
|
|
# Save the plot to the in-memory buffer as a PNG
|
|
plt.savefig(buffer, format="png")
|
|
buffer.seek(0)
|
|
|
|
# Create a Pillow Image object
|
|
image = Image.open(buffer)
|
|
images.append(image)
|
|
|
|
return (pil2tensor(images),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"SDT_PlotWaveForm": PlotWaveForm,
|
|
"SDT_PlotSpecgram": PlotSpecgram,
|
|
"SDT_PlotSpectrogram": PlotSpectrogram,
|
|
"SDT_PlotMelFilterBank": PlotMelFilterBank,
|
|
"SDT_PlotPitch": PlotPitch,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"SDT_PlotWaveForm": "Plot WaveForm",
|
|
"SDT_PlotSpecgram": "Plot Specgram",
|
|
"SDT_PlotSpectrogram": "Plot Spectrogram",
|
|
"SDT_PlotMelFilterBank": "Plot MelFilterBank",
|
|
"SDT_PlotPitch": "Plot Pitch",
|
|
}
|