tests, limiter
This commit is contained in:
c0ffymachyne
+250
-33
@@ -8,33 +8,40 @@ Version: 1.0.0
|
||||
Description:
|
||||
various plotting methods for debugging and visualization
|
||||
"""
|
||||
|
||||
import torch
|
||||
from io import BytesIO
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchaudio
|
||||
from PIL import Image
|
||||
import matplotlib.pyplot as plt
|
||||
from typing import Tuple, List
|
||||
from matplotlib.figure import Figure
|
||||
|
||||
|
||||
def _figure_to_image(figure, format="png", dpi=96):
|
||||
buf = BytesIO()
|
||||
figure.savefig(buf, format=format, dpi=96)
|
||||
plt.close(figure)
|
||||
def _figure_to_image(figure: Figure, dpi: int = 96) -> Image.Image:
|
||||
"""Convert a Matplotlib figure to a high-resolution RGB PIL Image."""
|
||||
figure.set_dpi(dpi)
|
||||
figure.canvas.draw()
|
||||
data = np.frombuffer(figure.canvas.tostring_argb(), dtype=np.uint8)
|
||||
width, height = figure.canvas.get_width_height()
|
||||
image = data.reshape((height, width, 4)) # ARGB format
|
||||
|
||||
buf.seek(0)
|
||||
waveform_image = Image.open(buf).convert("RGB")
|
||||
# Convert ARGB to RGB
|
||||
rgb_image = np.zeros((height, width, 3), dtype=np.uint8)
|
||||
rgb_image[:, :, 0] = image[:, :, 1] # Red
|
||||
rgb_image[:, :, 1] = image[:, :, 2] # Green
|
||||
rgb_image[:, :, 2] = image[:, :, 3] # Blue
|
||||
|
||||
return waveform_image
|
||||
return Image.fromarray(rgb_image)
|
||||
|
||||
|
||||
def get_wave(
|
||||
waveform: torch.Tensor,
|
||||
sample_rate: int,
|
||||
title: str = "Waveform",
|
||||
xlim=None,
|
||||
ylim=None,
|
||||
):
|
||||
waveform = waveform.numpy()
|
||||
xlim: int = 1000,
|
||||
ylim: int = 1000,
|
||||
) -> Image:
|
||||
waveform = waveform.cpu().numpy()
|
||||
|
||||
num_channels, num_frames = waveform.shape
|
||||
time_axis = torch.arange(0, num_frames) / sample_rate
|
||||
@@ -61,27 +68,237 @@ def get_wave(
|
||||
|
||||
|
||||
def get_spectogram(
|
||||
waveform: torch.Tensor, sample_rate: int, title="Spectrogram", xlim=None
|
||||
):
|
||||
waveform = waveform.numpy()
|
||||
num_channels, num_frames = waveform.shape
|
||||
time_axis = torch.arange(0, num_frames) / sample_rate
|
||||
waveform: torch.Tensor,
|
||||
sample_rate: int,
|
||||
n_fft: int = 4096,
|
||||
n_mels: int = 512,
|
||||
title: str = "Spectrogram",
|
||||
xlim: int = 8192,
|
||||
dpi: int = 96, # Set a high DPI for better image resolution
|
||||
) -> np.ndarray:
|
||||
"""Generate and plot a high-resolution spectrogram from a waveform."""
|
||||
|
||||
figure, axes = plt.subplots(num_channels, 1)
|
||||
if num_channels == 1:
|
||||
axes = [axes]
|
||||
for c in range(num_channels):
|
||||
axes[c].specgram(waveform[c], Fs=sample_rate)
|
||||
if num_channels > 1:
|
||||
axes[c].set_ylabel(f"Channel {c+1}")
|
||||
if xlim:
|
||||
axes[c].set_xlim(xlim)
|
||||
figure.suptitle(title)
|
||||
# Parameters for Mel Spectrogram
|
||||
win_length = n_fft // 2
|
||||
hop_length = n_fft // 4 # Smaller hop for better time resolution
|
||||
|
||||
waveform_image = _figure_to_image(figure)
|
||||
spectrogram_transform = torchaudio.transforms.MelSpectrogram(
|
||||
sample_rate=sample_rate,
|
||||
n_fft=n_fft,
|
||||
win_length=win_length,
|
||||
hop_length=hop_length,
|
||||
center=True,
|
||||
pad_mode="reflect",
|
||||
normalized=True,
|
||||
power=2.0, # Using power spectrogram
|
||||
norm="slaney",
|
||||
n_mels=n_mels,
|
||||
mel_scale="slaney",
|
||||
).to(waveform.device, dtype=waveform.dtype)
|
||||
|
||||
# Compute spectrogram
|
||||
mel_spectrogram = spectrogram_transform(waveform).cpu()
|
||||
|
||||
# Convert to decibel scale for better visualization
|
||||
spectrogram = torchaudio.transforms.AmplitudeToDB(top_db=80)(
|
||||
mel_spectrogram
|
||||
).numpy()
|
||||
|
||||
# Plot the spectrogram
|
||||
num_channels, _ = waveform.shape
|
||||
figure, axes = plt.subplots(
|
||||
num_channels, 1, figsize=(20, 10 * num_channels), squeeze=False, dpi=dpi
|
||||
)
|
||||
figure.suptitle(title, fontsize=16)
|
||||
|
||||
for i, ax in enumerate(axes[:, 0]): # Unpack axes
|
||||
ax.imshow(
|
||||
spectrogram[i],
|
||||
origin="lower",
|
||||
aspect="auto",
|
||||
extent=[0, xlim, 0, sample_rate / 2],
|
||||
cmap="magma",
|
||||
)
|
||||
ax.set_title(f"Channel {i + 1}", fontsize=14)
|
||||
ax.set_xlabel("Time (frames)", fontsize=12)
|
||||
ax.set_ylabel("Frequency (Hz)", fontsize=12)
|
||||
ax.tick_params(axis="both", which="major", labelsize=10)
|
||||
|
||||
plt.tight_layout(rect=[0, 0.03, 1, 0.95])
|
||||
|
||||
# Convert the figure to a high-resolution image
|
||||
waveform_image = _figure_to_image(figure, dpi=dpi)
|
||||
plt.close(figure) # Close the figure to free up memory
|
||||
return waveform_image
|
||||
|
||||
|
||||
def save_image(filepath: str, image: Image):
|
||||
image.save(filepath, format="JPEG", quality=95)
|
||||
def save_image(filepath: str, image: Image) -> Image:
|
||||
image.save(filepath, format="PNG", quality=95)
|
||||
|
||||
|
||||
def save_harmonic_spectrum(
|
||||
waveform: torch.Tensor,
|
||||
sample_rate: int,
|
||||
output_image: str,
|
||||
figsize: Tuple[int, int] = (12, 6),
|
||||
num_harmonics: int = 10,
|
||||
) -> None:
|
||||
"""
|
||||
Generate the harmonic spectrum of a waveform and save it as an image.
|
||||
|
||||
Parameters:
|
||||
waveform (torch.Tensor): Audio waveform tensor.
|
||||
sample_rate (int): Sample rate of the audio.
|
||||
output_image (str): Path to save the output image.
|
||||
figsize (tuple): Size of the output figure in inches.
|
||||
num_harmonics (int): Number of harmonics to calculate.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Ensure mono audio (combine channels if necessary)
|
||||
if waveform.size(0) > 1:
|
||||
waveform = torch.mean(waveform, dim=0, keepdim=True)
|
||||
|
||||
# Perform FFT to get the frequency domain
|
||||
fft = torch.fft.fft(waveform).to(device=waveform.device)
|
||||
magnitude = torch.abs(fft[0]) # Magnitude of the FFT
|
||||
frequencies = torch.fft.fftfreq(waveform.size(1), d=1 / sample_rate).to(
|
||||
device=waveform.device
|
||||
)
|
||||
|
||||
# Extract the fundamental frequency
|
||||
fundamental_idx = torch.argmax(magnitude[: len(magnitude) // 2])
|
||||
fundamental_freq = frequencies[fundamental_idx]
|
||||
|
||||
# Calculate harmonic frequencies
|
||||
harmonic_frequencies = [fundamental_freq * (i + 1) for i in range(num_harmonics)]
|
||||
harmonic_amplitudes = [
|
||||
magnitude[int(harmonic / sample_rate * len(magnitude))]
|
||||
for harmonic in harmonic_frequencies
|
||||
]
|
||||
harmonic_frequencies = torch.tensor(harmonic_frequencies).cpu().numpy()
|
||||
harmonic_amplitudes = torch.tensor(harmonic_amplitudes).cpu().numpy()
|
||||
|
||||
# Plot the harmonic spectrum
|
||||
plt.figure(figsize=figsize)
|
||||
plt.plot(harmonic_frequencies, harmonic_amplitudes, color="blue", linewidth=2)
|
||||
plt.title("Harmonic Spectrum")
|
||||
plt.xlabel("Frequency (Hz)")
|
||||
plt.ylabel("Amplitude")
|
||||
plt.grid(True)
|
||||
|
||||
# Save the image
|
||||
plt.tight_layout()
|
||||
plt.savefig(output_image, dpi=300)
|
||||
plt.close()
|
||||
|
||||
|
||||
PREDEFINED_COLORS = [
|
||||
"#1f77b4",
|
||||
"#ff7f0e",
|
||||
"#2ca02c",
|
||||
"#d62728",
|
||||
"#9467bd",
|
||||
"#8c564b",
|
||||
"#e377c2",
|
||||
"#7f7f7f",
|
||||
"#bcbd22",
|
||||
"#17becf",
|
||||
]
|
||||
|
||||
|
||||
def plot_multiple_harmonic_spectra(
|
||||
audio_data: List[Tuple[torch.Tensor, int, str]],
|
||||
output_image: str,
|
||||
figsize: Tuple[int, int] = (12, 6),
|
||||
num_harmonics: int = 16,
|
||||
upper_bound: int = 140000,
|
||||
title: str = "Harmonic Spectrum",
|
||||
) -> None:
|
||||
"""
|
||||
Optimized: Plot high-resolution harmonic spectra for multiple audio waveforms.
|
||||
|
||||
Parameters:
|
||||
audio_data (list): List of (waveform, sample_rate, label) tuples.
|
||||
output_image (str): Path to save the combined output image.
|
||||
figsize (tuple): Size of the figure.
|
||||
num_harmonics (int): Number of harmonics to calculate.
|
||||
upper_bound (int): Maximum frequency to display.
|
||||
title (str): Title of the plot.
|
||||
"""
|
||||
plt.figure(figsize=figsize)
|
||||
|
||||
for idx, (waveform, sample_rate, label) in enumerate(audio_data):
|
||||
# Ensure mono audio
|
||||
if waveform.size(0) > 1:
|
||||
waveform = torch.mean(waveform, dim=0)
|
||||
|
||||
# Move waveform to GPU for efficient computation
|
||||
device = (
|
||||
waveform.device
|
||||
if waveform.is_cuda
|
||||
else "cuda" if torch.cuda.is_available() else "cpu"
|
||||
)
|
||||
waveform = waveform.to(device)
|
||||
|
||||
# High-resolution FFT (Zero-padding for better frequency resolution)
|
||||
n_fft = 1 * waveform.size(0) # Zero-padding factor of 4
|
||||
fft = torch.fft.fft(waveform, n=n_fft).to(device=waveform.device)
|
||||
magnitude = torch.abs(fft[: n_fft // 2])
|
||||
frequencies = torch.fft.fftfreq(n_fft, d=1 / sample_rate)[: n_fft // 2]
|
||||
|
||||
# Find fundamental frequency
|
||||
fundamental_idx = torch.argmax(
|
||||
magnitude[: len(magnitude) // 4]
|
||||
) # Search in the first quarter
|
||||
fundamental_freq = frequencies[fundamental_idx]
|
||||
|
||||
# Precompute harmonic frequencies
|
||||
harmonic_freqs = fundamental_freq * torch.arange(
|
||||
1, num_harmonics + 1, device=device
|
||||
)
|
||||
harmonic_indices = (harmonic_freqs / (sample_rate / n_fft)).long()
|
||||
harmonic_amplitudes = magnitude[harmonic_indices].cpu().numpy()
|
||||
|
||||
# Limit harmonic frequencies to upper bound
|
||||
harmonic_freqs = harmonic_freqs[harmonic_freqs <= upper_bound].cpu().numpy()
|
||||
harmonic_amplitudes = harmonic_amplitudes[: len(harmonic_freqs)]
|
||||
|
||||
# Convert data to CPU for plotting
|
||||
frequencies_cpu = frequencies.cpu().numpy()
|
||||
magnitude_cpu = magnitude.cpu().numpy()
|
||||
|
||||
# Plot the full spectrum
|
||||
color = PREDEFINED_COLORS[idx % len(PREDEFINED_COLORS)]
|
||||
plt.plot(
|
||||
frequencies_cpu,
|
||||
magnitude_cpu,
|
||||
color=color,
|
||||
alpha=0.4,
|
||||
label=f"Full Spectrum ({label})",
|
||||
)
|
||||
|
||||
# Overlay harmonic peaks
|
||||
plt.vlines(
|
||||
harmonic_freqs,
|
||||
ymin=0,
|
||||
ymax=harmonic_amplitudes,
|
||||
color=color,
|
||||
linewidth=1.5,
|
||||
linestyle="--",
|
||||
label=f"Harmonics ({label})",
|
||||
)
|
||||
|
||||
plt.title(title)
|
||||
plt.xscale("log")
|
||||
plt.xlim(20, upper_bound)
|
||||
plt.xlabel("Frequency (Hz)")
|
||||
plt.ylabel("Amplitude")
|
||||
plt.legend()
|
||||
plt.grid(True, which="both", linestyle="--", linewidth=0.5)
|
||||
|
||||
# Save and close the figure
|
||||
plt.tight_layout()
|
||||
plt.savefig(output_image, dpi=150)
|
||||
plt.close()
|
||||
|
||||
Reference in New Issue
Block a user