import os

import wave

from time import sleep

from typing import List

from librosa.display import specshow, waveshow

from matplotlib import pyplot as plt

from pyaudio import PyAudio

from rainbowgram.rain2graph import rain2graph

from rainbowgram.wave_rain import wave2rain

from tqdm import tqdm

from autoencoder import VAE

from config import *

from preprocess import MinMaxNormaliser

from soundgenerator import convert_spectrograms_to_audio, play_sounds

from train import load_dataset

chunk = cfg.SAMPLE_RATE

def play_all_original_sounds():

directory = "../datasets/one_shot_percussive_sounds/audio"

p = PyAudio()

stream = p.open(format=2,

channels=1,

rate=cfg.SAMPLE_RATE,

output=True)

for file in tqdm(os.listdir(directory)):

filename = os.fsdecode(file)

if filename.endswith(".wav"):

f = wave.open(os.path.join(directory, filename), "rb")

data = np.frombuffer(f.readframes(chunk), dtype=np.int16)

data = np.pad(data, (0, chunk - f.getnframes()), 'constant')

stream.write(np.frombuffer(data, dtype=np.int16))

else:

continue

stream.stop_stream()

stream.close()

def play_all_spectrogrammed_sounds():

x_train, _ = load_dataset(num_samples=100)

play_sounds(x_train, is_spec=True)

def play_all_reconstructed_sounds():

x_train, file_paths = load_dataset(num_samples=200)

vae = VAE.load("model_220621")

play_sounds(

vae.model.predict(x_train),

is_spec=True,

# print_names=file_paths

)

def show_rain(filename: str):

f = wave.open(filename, "rb")

data_original = np.frombuffer(f.readframes(np.inf), dtype=np.int16)

data_original = np.pad(data_original, (0, chunk - len(data_original)), 'constant')

normalization_factor = np.max(np.abs(data_original)).astype(np.float64)

data_original = data_original / normalization_factor

rain_original = wave2rain(data_original, sr=cfg.SAMPLE_RATE, n_fft=cfg.N_FFT, stride=cfg.HOP_LENGTH)

rain2graph(rain_original)

plt.show(dpi=600)

def pretty_plot_spectrogram(

spectrogram: np.ndarray,

signal: Optional[np.ndarray] = None,

title: str = ''):

if signal is not None:

# plt.figure()

fig, axs = plt.subplots(2, 1, constrained_layout=True)

fig.suptitle(title, fontsize=16)

waveshow(signal / cfg.NORMALIZATION_FACTOR, sr=cfg.SAMPLE_RATE, ax=axs[0])

# divider = make_axes_locatable(axs[0])

# cax = divider.append_axes("right", size="5%", pad=.05)

axs[0].set_xticks(np.linspace(0, 1, 11))

axs[0].set_xlim([0, 1])

axs[0].set_title('Signal')

axs[0].set_xlabel('Time [s]')

axs[0].set_ylabel('Amplitude')

img = specshow(spectrogram[..., 0], ax=axs[1],

x_axis='time', y_axis='log', sr=cfg.SAMPLE_RATE, n_fft=cfg.N_FFT, hop_length=cfg.HOP_LENGTH)

cax = axs[1].inset_axes([1.01, 0.0, 0.02, 1.0])

cbar = plt.colorbar(img, format="%+2.f dB", ax=axs[1], cax=cax)

cbar.ax.tick_params(labelsize=6)

axs[1].set_xticks(np.linspace(0, 1, 11))

axs[1].set_xlim([0, 1])

axs[1].set_title('Spectrogram')

axs[1].set_xlabel('Time [s]')

axs[1].set_ylabel('Frequency [Hz]')

# plt.tight_layout()

plt.show()

else:

fig, ax = plt.subplots()

ax.set(title=title)

img = specshow(spectrogram[..., 0], ax=ax,

x_axis='time', y_axis='log', sr=cfg.SAMPLE_RATE, n_fft=cfg.N_FFT, hop_length=cfg.HOP_LENGTH)

ax.set_xticks(np.linspace(0, 1, 11))

ax.set_xlim([0, 1])

ax.set_xlabel('Time [s]')

ax.set_ylabel('Frequency [Hz]')

fig.colorbar(img, ax=ax, format="%+2.f dB")

plt.gcf().set_dpi(300)

plt.show()

def show_reconstruction_examples():

num_samples = 1

show = True

play = False

vae = VAE.load("model_220623")

m, _ = load_dataset(num_samples=num_samples)

z = vae.model.predict(m)

if play:

play_sounds(convert_spectrograms_to_audio(m), is_spec=False)

sleep(1)

play_sounds(convert_spectrograms_to_audio(z), is_spec=False)

if show:

m_signals = convert_spectrograms_to_audio(m)

z_signals = convert_spectrograms_to_audio(z)

m_specs = MinMaxNormaliser.denormalise(m)

z_specs = MinMaxNormaliser.denormalise(z)

for i in range(num_samples):

pretty_plot_spectrogram(m_specs[i], m_signals[i], title=f"Sample - {i + 1}(/{num_samples})")

pretty_plot_spectrogram(z_specs[i], z_signals[i], title=f"Sample reconstruction - {i + 1}(/{num_samples})")

# Zero response

pretty_plot_spectrogram(vae.model.predict(np.zeros((1, 128, 128, 1)))[0], title=f"Zeros reconstruction")

def show_interpolation(

interpolated_specs: np.ndarray,

interpolated_signals: List[np.ndarray],

title: str) \

-> None:

assert len(interpolated_specs) == len(interpolated_signals) == 9, \

"It is hardcoded having 9 total steps of interpolation including edges"

pretty_plot_spectrogram(interpolated_specs[0], interpolated_signals[0], title=f"Interpolation {title} - source")

pretty_plot_spectrogram(interpolated_specs[2], interpolated_signals[2], title=f"Interpolation {title} - 25%")

pretty_plot_spectrogram(interpolated_specs[4], interpolated_signals[4], title=f"Interpolation {title} - 50%")

pretty_plot_spectrogram(interpolated_specs[6], interpolated_signals[6], title=f"Interpolation {title} - 75%")

pretty_plot_spectrogram(interpolated_specs[8], interpolated_signals[8], title=f"Interpolation {title} - target")

def play_interpolation(

interpolated_specs: np.ndarray,

save_dir: Optional[str] = None) \

-> None:

output_specs = np.stack([

interpolated_specs[0], interpolated_specs[0], cfg.EMPTY_SPEC,

interpolated_specs[-1], interpolated_specs[-1], cfg.EMPTY_SPEC,

*interpolated_specs

])

play_sounds(output_specs, save_dir=save_dir, save_joined=True)