reorg

core/plotting.py

@@ -0,0 +1,76 @@
+#!/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)