| from os import path |
|
|
| import librosa as rosa |
| import matplotlib |
| import matplotlib.pyplot as plt |
| import numpy as np |
| from pytorch_lightning.loggers import TensorBoardLogger |
| from pytorch_lightning.utilities import rank_zero_only |
|
|
| from utils.stft import STFTMag |
|
|
| matplotlib.use('Agg') |
|
|
|
|
| class TensorBoardLoggerExpanded(TensorBoardLogger): |
| def __init__(self, sr=16000): |
| super().__init__(save_dir='lightning_logs', default_hp_metric=False, name='') |
| self.sr = sr |
| self.stftmag = STFTMag() |
|
|
| def fig2np(self, fig): |
| data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='') |
| data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,)) |
| return data |
|
|
| def plot_spectrogram_to_numpy(self, y, y_low, y_recon, step): |
| name_list = ['y', 'y_low', 'y_recon'] |
| fig = plt.figure(figsize=(9, 15)) |
| fig.suptitle(f'Epoch_{step}') |
| for i, yy in enumerate([y, y_low, y_recon]): |
| if yy.dim() == 1: |
| yy = self.stftmag(yy) |
| ax = plt.subplot(3, 1, i + 1) |
| ax.set_title(name_list[i]) |
| plt.imshow(rosa.amplitude_to_db(yy.numpy(), |
| ref=np.max, top_db=80.), |
| |
| vmax=0., |
| aspect='auto', |
| origin='lower', |
| interpolation='none') |
| plt.colorbar() |
| plt.xlabel('Frames') |
| plt.ylabel('Channels') |
| plt.tight_layout() |
|
|
| fig.canvas.draw() |
| data = self.fig2np(fig) |
|
|
| plt.close() |
| return data |
|
|
| @rank_zero_only |
| def log_spectrogram(self, y, y_low, y_recon, epoch): |
| y, y_low, y_recon = y.detach().cpu(), y_low.detach().cpu(), y_recon.detach().cpu() |
| spec_img = self.plot_spectrogram_to_numpy(y, y_low, y_recon, epoch) |
| self.experiment.add_image(path.join(self.save_dir, 'result'), |
| spec_img, |
| epoch, |
| dataformats='HWC') |
| self.experiment.flush() |
| return |
|
|
| @rank_zero_only |
| def log_audio(self, y, y_low, y_recon, epoch): |
| y, y_low, y_recon = y.detach().cpu(), y_low.detach().cpu(), y_recon.detach().cpu(), |
| name_list = ['y', 'y_low', 'y_recon'] |
| for n, yy in zip(name_list, [y, y_low, y_recon]): |
| self.experiment.add_audio(n, yy, epoch, self.sr) |
| self.experiment.flush() |
| return |
|
|
