127 lines
4.7 KiB
Python
127 lines
4.7 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
Author: C0ffymachyne
|
|
License: GPLv3
|
|
Version: 1.0.0
|
|
|
|
Description:
|
|
Spectogram image node
|
|
"""
|
|
|
|
import torch
|
|
|
|
from typing import List, Dict
|
|
from PIL import Image
|
|
|
|
import numpy as np
|
|
import torchaudio
|
|
|
|
import matplotlib.pyplot as plt
|
|
|
|
from ..core.io import audio_to_comfy_3d, audio_from_comfy_3d, audio_from_comfy_2d
|
|
from ..core.loudness import lufs_normalization, get_loudness
|
|
|
|
class SignalProcessingSpectrogram:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
cmaps = ['viridis', 'plasma', 'inferno', 'magma', 'cividis']
|
|
return {
|
|
"required": {
|
|
"audio_input": ("AUDIO",),
|
|
"color_map": (cmaps,),
|
|
},
|
|
"optional": {
|
|
"n_fft": ("INT", {"default": 4096, "min": 512, "max": 8192, "step": 256}),
|
|
"hop_length": ("INT", {"default": 128, "min": 64, "max": 4096, "step": 128}),
|
|
"n_mels": ("INT", {"default": 512, "min": 32, "max": 2048, "step": 32}),
|
|
"top_db": ("FLOAT", {"default": 80.0, "min": 10.0, "max": 100.0, "step": 5.0}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("spectrogram_image",)
|
|
CATEGORY = "Audio Processing"
|
|
FUNCTION = "process"
|
|
|
|
def process(self, audio_input, color_map, n_fft=2048, hop_length=512, n_mels=128, top_db=80.0):
|
|
waveform = audio_input.get('waveform') # [channels, samples] or [batch, channels, samples]
|
|
sample_rate = audio_input.get('sample_rate')
|
|
|
|
#waveform, sample_rate = audio_from_comfy_2d(audio_input)
|
|
|
|
# Convert to mono by averaging channels
|
|
if waveform.ndim == 3:
|
|
# [batch, channels, samples]
|
|
waveform = waveform.mean(dim=1, keepdim=True) # [batch, 1, samples]
|
|
waveform = waveform.squeeze(0) # [1, samples]
|
|
elif waveform.ndim == 2:
|
|
# [channels, samples]
|
|
if waveform.shape[0] > 1:
|
|
waveform = waveform.mean(dim=0, keepdim=True) # [1, samples]
|
|
else:
|
|
waveform = waveform.unsqueeze(0) # [1, samples]
|
|
elif waveform.ndim == 1:
|
|
# [samples]
|
|
waveform = waveform.unsqueeze(0) # [1, samples]
|
|
else:
|
|
raise ValueError(f"Unsupported waveform shape: {waveform.shape}")
|
|
|
|
# Generate Mel Spectrogram
|
|
spectrogram_transform = torchaudio.transforms.MelSpectrogram(
|
|
sample_rate=sample_rate,
|
|
n_fft=n_fft,
|
|
hop_length=hop_length,
|
|
n_mels=n_mels,
|
|
power=2.0,
|
|
norm='slaney',
|
|
mel_scale='htk',
|
|
).to(waveform.device, dtype = waveform.dtype)
|
|
spectrogram = spectrogram_transform(waveform) # [1, n_mels, time_frames]
|
|
|
|
# Convert to decibel scale
|
|
amplitude_to_db = torchaudio.transforms.AmplitudeToDB(top_db=top_db)
|
|
spectrogram_db = amplitude_to_db(spectrogram) # [1, n_mels, time_frames]
|
|
|
|
# Convert to numpy
|
|
spectrogram_db = spectrogram_db.squeeze().detach().cpu().numpy() # [n_mels, time_frames]
|
|
|
|
# Clip spectrogram to a range for better contrast
|
|
spectrogram_db = np.clip(spectrogram_db, -top_db, 0.0)
|
|
|
|
# Normalize spectrogram to [0,1]
|
|
spectrogram_normalized = (spectrogram_db + top_db) / top_db # [0,1]
|
|
|
|
# Apply a colormap (e.g., 'inferno') using matplotlib
|
|
cmap = plt.get_cmap(color_map)
|
|
spectrogram_colored = cmap(spectrogram_normalized) # [n_mels, time_frames, 4] RGBA
|
|
|
|
# Convert to RGB by removing alpha channel
|
|
spectrogram_rgb = (spectrogram_colored[:, :, :3] * 255).astype(np.uint8) # [n_mels, time_frames, 3]
|
|
spectrogram_rgb = np.squeeze(spectrogram_rgb)
|
|
|
|
# Check the shape and adjust if necessary
|
|
if len(spectrogram_rgb.shape) == 3 and spectrogram_rgb.shape[-1] == 3:
|
|
# Ensure the array is in uint8 format (0-255 range)
|
|
spectrogram_rgb = np.clip(spectrogram_rgb, 0, 255).astype(np.uint8)
|
|
else:
|
|
raise ValueError(f"Unexpected spectrogram shape: {spectrogram_rgb.shape}")
|
|
|
|
# Convert to RGB image
|
|
spectrogram_image = Image.fromarray(spectrogram_rgb).convert("RGB")
|
|
|
|
# Optionally resize for better resolution
|
|
spectrogram_image = spectrogram_image.resize(
|
|
(spectrogram_image.width * 2, spectrogram_image.height * 2),
|
|
Image.BILINEAR
|
|
)
|
|
|
|
# Convert to numpy array and normalize to [0,1]
|
|
image_np = np.array(spectrogram_image).astype(np.float32) / 255.0 # [H, W, 3]
|
|
|
|
# Convert to torch tensor and add batch dimension
|
|
#image_tensor = torch.from_numpy(image_np).permute(2, 0, 1).unsqueeze(0) # [1, 3, H, W]
|
|
|
|
image = torch.from_numpy(image_np)[None,]
|
|
|
|
return (image, ) |