justyoung/DiffSinger

1import matplotlib.pyplot as plt2import numpy as np3import torch4 5LINE_COLORS = ['w', 'r', 'y', 'cyan', 'm', 'b', 'lime']6 7 8def spec_to_figure(spec, vmin=None, vmax=None):9    if isinstance(spec, torch.Tensor):10        spec = spec.cpu().numpy()11    fig = plt.figure(figsize=(12, 6))12    plt.pcolor(spec.T, vmin=vmin, vmax=vmax)13    return fig14 15 16def spec_f0_to_figure(spec, f0s, figsize=None):17    max_y = spec.shape[1]18    if isinstance(spec, torch.Tensor):19        spec = spec.detach().cpu().numpy()20        f0s = {k: f0.detach().cpu().numpy() for k, f0 in f0s.items()}21    f0s = {k: f0 / 10 for k, f0 in f0s.items()}22    fig = plt.figure(figsize=(12, 6) if figsize is None else figsize)23    plt.pcolor(spec.T)24    for i, (k, f0) in enumerate(f0s.items()):25        plt.plot(f0.clip(0, max_y), label=k, c=LINE_COLORS[i], linewidth=1, alpha=0.8)26    plt.legend()27    return fig28 29 30def dur_to_figure(dur_gt, dur_pred, txt):31    dur_gt = dur_gt.long().cpu().numpy()32    dur_pred = dur_pred.long().cpu().numpy()33    dur_gt = np.cumsum(dur_gt)34    dur_pred = np.cumsum(dur_pred)35    fig = plt.figure(figsize=(12, 6))36    for i in range(len(dur_gt)):37        shift = (i % 8) + 138        plt.text(dur_gt[i], shift, txt[i])39        plt.text(dur_pred[i], 10 + shift, txt[i])40        plt.vlines(dur_gt[i], 0, 10, colors='b')  # blue is gt41        plt.vlines(dur_pred[i], 10, 20, colors='r')  # red is pred42    return fig43 44 45def f0_to_figure(f0_gt, f0_cwt=None, f0_pred=None):46    fig = plt.figure()47    f0_gt = f0_gt.cpu().numpy()48    plt.plot(f0_gt, color='r', label='gt')49    if f0_cwt is not None:50        f0_cwt = f0_cwt.cpu().numpy()51        plt.plot(f0_cwt, color='b', label='cwt')52    if f0_pred is not None:53        f0_pred = f0_pred.cpu().numpy()54        plt.plot(f0_pred, color='green', label='pred')55    plt.legend()56    return fig57