76 lines
2.0 KiB
Python
76 lines
2.0 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 torch
|
|
from io import BytesIO
|
|
from PIL import Image
|
|
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
|
|
def _figure_to_image(figure,format='png',dpi=96):
|
|
buf = BytesIO()
|
|
figure.savefig(buf, format=format, dpi=96)
|
|
plt.close(figure)
|
|
|
|
buf.seek(0)
|
|
waveform_image = Image.open(buf).convert("RGB")
|
|
|
|
return waveform_image
|
|
|
|
def get_wave(waveform : torch.Tensor, sample_rate : int , title : str ="Waveform", xlim=None, ylim=None):
|
|
waveform = waveform.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 , title="Spectrogram", xlim=None):
|
|
waveform = waveform.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].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)
|
|
|
|
waveform_image = _figure_to_image(figure)
|
|
|
|
return waveform_image
|
|
|
|
def save_image(filepath : str, image : Image):
|
|
image.save(filepath, format="JPEG", quality=95) |