aboutsummaryrefslogtreecommitdiff
path: root/src/metrics.py
blob: 08eb3cd1f561f30c62a95f5b59a4764880a92d37 (plain) (blame)
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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
import numpy as np
import scipy


def fid(reals: np.ndarray, fakes: np.ndarray) -> float:
    """
    Calculate Frechet Inception Distance (FID) score.

    Args:
        reals: Real images [N, C, H, W]
        fakes: Generated images [N, C, H, W]

    Returns:
        FID score (lower is better)
    """
    reals = reals.reshape(reals.shape[0], -1)
    fakes = fakes.reshape(fakes.shape[0], -1)

    mu_real = np.mean(reals, axis=0)
    mu_fake = np.mean(fakes, axis=0)
    sigma_real = np.cov(reals, rowvar=False)
    sigma_fake = np.cov(fakes, rowvar=False)

    diff = mu_real - mu_fake
    covmean, _ = scipy.linalg.sqrtm(sigma_real.dot(sigma_fake), disp=False)

    if not np.isfinite(covmean).all():
        eps = 1e-6
        offset = np.eye(sigma_real.shape[0]) * eps
        covmean = scipy.linalg.sqrtm((sigma_real + offset).dot(sigma_fake + offset))

    if np.iscomplexobj(covmean):
        covmean = covmean.real

    return diff @ diff + np.trace(sigma_real) + np.trace(sigma_fake) - 2 * np.trace(covmean)


def kl_divergence(reals: np.ndarray, fakes: np.ndarray) -> float:
    """
    Calculate KL divergence between real and fake image distributions.

    Args:
        reals: Real images [N, C, H, W]
        fakes: Generated images [N, C, H, W]

    Returns:
        KL divergence value
    """
    reals = reals.transpose(1, 0, 2, 3).reshape(reals.shape[1], -1)
    fakes = fakes.transpose(1, 0, 2, 3).reshape(fakes.shape[1], -1)

    hist_real = np.apply_along_axis(
        lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, reals
    )
    hist_fake = np.apply_along_axis(
        lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, fakes
    )

    # Add smoothing
    hist_real = hist_real + 1
    hist_fake = hist_fake + 1

    hist_real = hist_real / np.sum(hist_real)
    hist_fake = hist_fake / np.sum(hist_fake)

    return np.mean(np.log(hist_real / hist_fake))


def jsd(reals: np.ndarray, fakes: np.ndarray) -> float:
    """
    Calculate Jensen-Shannon divergence between real and fake image distributions.

    Args:
        reals: Real images [N, C, H, W]
        fakes: Generated images [N, C, H, W]

    Returns:
        JSD value
    """
    reals = reals.transpose(1, 0, 2, 3).reshape(reals.shape[1], -1)
    fakes = fakes.transpose(1, 0, 2, 3).reshape(fakes.shape[1], -1)

    hist_real = np.apply_along_axis(
        lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, reals
    )
    hist_fake = np.apply_along_axis(
        lambda a: np.histogram(a, bins=40, range=(-1, 1))[0], 1, fakes
    )

    hist_real = hist_real + 1
    hist_fake = hist_fake + 1

    hist_real = hist_real / np.sum(hist_real)
    hist_fake = hist_fake / np.sum(hist_fake)

    hist_avg = (hist_real + hist_fake) / 2

    return 0.5 * (np.mean(np.log(hist_real / hist_avg)) + np.mean(np.log(hist_fake / hist_avg)))