92 lines
3.3 KiB
Python
92 lines
3.3 KiB
Python
# Copyright (C) 2024 Reece H. Dunn. SPDX-License-Identifier: GPL-3
|
|
|
|
import torch
|
|
import torchaudio.functional as F
|
|
|
|
|
|
class LinearFilterBank:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"spectrogram": ("SPECT",)},
|
|
"optional": {"f_min": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 96000.0, "step": 1.0}),
|
|
"f_max": ("FLOAT", {"default": 44100.0, "min": 0.0, "max": 96000.0, "step": 1.0}),
|
|
"n_filter": ("INT", {"default": 128, "min": 1, "max": 2**32})}}
|
|
|
|
CATEGORY = "audio processing/filter bank"
|
|
|
|
RETURN_TYPES = ("FILTER_BANK", )
|
|
FUNCTION = "fbank"
|
|
|
|
def fbank(self,
|
|
spectrogram,
|
|
f_min=0.0,
|
|
f_max=44100.0,
|
|
n_filter=128):
|
|
n_freqs = spectrogram["n_fft"] // 2 + 1
|
|
fbank = F.linear_fbanks(n_freqs=n_freqs,
|
|
f_min=f_min,
|
|
f_max=f_max,
|
|
n_filter=n_filter,
|
|
sample_rate=spectrogram["sample_rate"])
|
|
return ({"fbank": fbank,
|
|
"sample_rate": spectrogram["sample_rate"],
|
|
"n_freqs": n_freqs,
|
|
"f_min": f_min,
|
|
"f_max": f_max,
|
|
"n_bins": n_filter}, )
|
|
|
|
|
|
class MelScaleFilterBank:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"spectrogram": ("SPECT",)},
|
|
"optional": {"f_min": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 96000.0, "step": 1.0}),
|
|
"f_max": ("FLOAT", {"default": 44100.0, "min": 0.0, "max": 96000.0, "step": 1.0}),
|
|
"n_mels": ("INT", {"default": 128, "min": 1, "max": 2**32})}}
|
|
|
|
CATEGORY = "audio processing/filter bank"
|
|
|
|
RETURN_TYPES = ("FILTER_BANK", )
|
|
FUNCTION = "fbank"
|
|
|
|
def fbank(self,
|
|
spectrogram,
|
|
f_min=0.0,
|
|
f_max=44100.0,
|
|
n_mels=128):
|
|
n_freqs = spectrogram["n_fft"] // 2 + 1
|
|
fbank = F.melscale_fbanks(n_freqs=n_freqs,
|
|
f_min=f_min,
|
|
f_max=f_max,
|
|
n_mels=n_mels,
|
|
sample_rate=spectrogram["sample_rate"])
|
|
return ({"fbank": fbank,
|
|
"sample_rate": spectrogram["sample_rate"],
|
|
"n_freqs": n_freqs,
|
|
"f_min": f_min,
|
|
"f_max": f_max,
|
|
"n_bins": n_mels}, )
|
|
|
|
|
|
class ApplyFilterBank:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": {"spectrogram": ("SPECT",),
|
|
"fbank": ("FILTER_BANK",)}}
|
|
|
|
CATEGORY = "audio processing/filter bank"
|
|
|
|
RETURN_TYPES = ("SPECT", )
|
|
FUNCTION = "filter"
|
|
|
|
def filter(self, spectrogram, fbank):
|
|
spect = spectrogram["spectrogram"]
|
|
fb = fbank["fbank"]
|
|
filtered = torch.matmul(spect.transpose(-1, -2), fb).transpose(-1, -2)
|
|
return ({"spectrogram": filtered,
|
|
"sample_rate": spectrogram["sample_rate"],
|
|
"stype": spectrogram["stype"],
|
|
"n_fft": spectrogram["n_fft"],
|
|
"win_length": spectrogram["win_length"],
|
|
"hop_length": spectrogram["hop_length"]}, )
|