File size: 1,658 Bytes
32b2aaa
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d2b7e94
32b2aaa
 
 
d2b7e94
32b2aaa
 
 
d2b7e94
32b2aaa
 
 
d2b7e94
32b2aaa
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
from typing import Callable, TypeVar, overload

import matplotlib.pyplot as plt
import numpy as np


def save_mels(path, *, targ_mel, pred_mel, cond_mel):
    n = 3 if cond_mel is None else 4

    plt.figure(figsize=(10, n * 4))

    i = 1

    plt.subplot(n, 1, i)
    plt.imshow(pred_mel, origin="lower", interpolation="none")
    plt.title(f"Pred mel {pred_mel.shape}")
    i += 1

    plt.subplot(n, 1, i)
    plt.imshow(targ_mel, origin="lower", interpolation="none")
    plt.title(f"GT mel {targ_mel.shape}")
    i += 1

    plt.subplot(n, 1, i)
    pred_mel = pred_mel[:, : targ_mel.shape[1]]
    targ_mel = targ_mel[:, : pred_mel.shape[1]]
    plt.imshow(np.abs(pred_mel - targ_mel), origin="lower", interpolation="none")
    plt.title(f"Diff mel {pred_mel.shape}, mse={np.mean((pred_mel - targ_mel)**2):.4f}")
    i += 1

    if cond_mel is not None:
        plt.subplot(n, 1, i)
        plt.imshow(cond_mel, origin="lower", interpolation="none")
        plt.title(f"Cond mel {cond_mel.shape}")
        i += 1

    plt.savefig(path, dpi=480)
    plt.close()


T = TypeVar("T")


@overload
def tree_map(fn: Callable, x: list[T]) -> list[T]: ...


@overload
def tree_map(fn: Callable, x: tuple[T]) -> tuple[T]: ...


@overload
def tree_map(fn: Callable, x: dict[str, T]) -> dict[str, T]: ...


@overload
def tree_map(fn: Callable, x: T) -> T: ...


def tree_map(fn: Callable, x):
    if isinstance(x, list):
        x = [tree_map(fn, xi) for xi in x]
    elif isinstance(x, tuple):
        x = (tree_map(fn, xi) for xi in x)
    elif isinstance(x, dict):
        x = {k: tree_map(fn, v) for k, v in x.items()}
    else:
        x = fn(x)
    return x