implement ation JS divergence

2 files changed, 30 insertions, 8 deletions

diff --git a/main.py b/main.py
index 81f832d..7c7daa1 100644
--- a/main.py
+++ b/main.py
@@ -423,5 +423,11 @@ if __name__ == '__main__':
         kl_score = kl(reals, fakes)
         print(f"KL divergence: {kl_score}")
 
+        rkl_score = kl(fakes, reals)
+        print(f"Reverse KL divergence: {rkl_score}")
+
+        jsd_score = jsd(reals, fakes)
+        print(f"Jensen-Shannon divergence: {jsd_score}")
+
 
         # plt.show()
diff --git a/utils.py b/utils.py
index 4e680a8..509bd11 100644
--- a/utils.py
+++ b/utils.py
@@ -47,13 +47,27 @@ def kl(reals, fakes):
     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)
 
-    plt.figure()
-    colors = ['#ff0000', '#00ff00', '#0000ff']
-    for i in range(reals.shape[0]):
-        plt.plot(hist_real[i], label='real', color=colors[i])
-        plt.plot(hist_fake[i], label='fake', color=colors[i], linestyle='dashed')
-    plt.legend()
-    plt.savefig('hist.tmp.png')
+    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, fakes):
+    """Jensen-Shannon divergence calculation.
+
+    Args:
+        reals (numpy.array): Real images.
+        fakes (numpy.array): Fake images.
+    """
+    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
@@ -61,4 +75,6 @@ def kl(reals, fakes):
     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))
\ No newline at end of file
+    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)))
\ No newline at end of file