305 lines
9.2 KiB
Python
305 lines
9.2 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
Author: C0ffymachyne
|
|
License: GPLv3
|
|
Version: 1.0.0
|
|
|
|
Description:
|
|
various plotting methods for debugging and visualization
|
|
"""
|
|
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: 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
|
|
|
|
# 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 Image.fromarray(rgb_image)
|
|
|
|
|
|
def get_wave(
|
|
waveform: torch.Tensor,
|
|
sample_rate: int,
|
|
title: str = "Waveform",
|
|
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
|
|
|
|
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_ylabel(f"Channel {c+1}")
|
|
if xlim:
|
|
axes[c].set_xlim(xlim)
|
|
if ylim:
|
|
axes[c].set_ylim(ylim)
|
|
|
|
figure.suptitle(title)
|
|
|
|
waveform_image = _figure_to_image(figure)
|
|
|
|
return waveform_image
|
|
|
|
|
|
def get_spectogram(
|
|
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."""
|
|
|
|
# Parameters for Mel Spectrogram
|
|
win_length = n_fft // 2
|
|
hop_length = n_fft // 4 # Smaller hop for better time resolution
|
|
|
|
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:
|
|
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()
|