Spaces:
Running on Zero
Running on Zero
File size: 1,472 Bytes
2e1dc7f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | from pathlib import Path
from typing import Optional
from matplotlib import pyplot as plt
import torch
from torch import Tensor, norm
import matplotlib
import config as cfg
from utils.audio import normalize
def plot_waveforms(*waveforms: Tensor,
savepath: Optional[Path] = None,
**kwargs):
if savepath is not None:
b = matplotlib.get_backend()
matplotlib.use("agg")
t = torch.linspace(0, waveforms[0].shape[-1],
waveforms[0].shape[-1]) # 5 seconds of audio
# audio_tensor = torch.sin(2 * np.pi * freq * t) # Generate sinewave
# Plot the waveform
plt.style.use(cfg.ROOT / "tkol.mplstyle")
colors = [c["color"] for c in list(plt.rcParams["axes.prop_cycle"])]
plt.figure(**kwargs)
for i, wave in enumerate(waveforms):
wave = wave.squeeze()
wave = normalize(wave, -1, 1)
# If your audio is stereo (2 channels), you can average over channels, or just plot one
if wave.ndimension() > 1:
wave = wave.mean(dim=0) # Take the mean over channels if stereo
plt.plot(t, wave, color=colors[i])
plt.title("Waveform")
plt.xlabel("Time")
plt.ylabel("Amplitude")
plt.gca().set_xticks([])
plt.gca().set_yticks([])
for s in plt.gca().spines.values():
s.set_visible((False))
if savepath is not None:
plt.savefig(savepath)
matplotlib.use(b)
else:
plt.show()
|