rewrote latent space tests

1 files changed, 139 insertions, 14 deletions

diff --git a/test/latent_space_test.py b/test/latent_space_test.py
index 0655923..d05eb78 100644
--- a/test/latent_space_test.py
+++ b/test/latent_space_test.py
@@ -1,9 +1,10 @@
 import unittest
 import numpy as np
 from scipy.stats import truncnorm
+import context
 from gantools import latent_space
-from gantools import biggan
 import PIL.Image
+import math
 
 def create_random_keyframe(n_vector, n_label):
     truncation = (0.9 - 0.1)*np.random.random() + 0.1
@@ -29,21 +30,145 @@ def save_ims(ims):
         i += 1
 ####
 
+def compare_float_arrays_2d(target_seq, actual_seq):
+    for target, actual in zip(target_seq, actual_seq):
+        for ti, ai in zip(target, actual):
+            assert math.isclose(ti, ai), 'target: %s; actual %s' % (str(target), str(actual))
+
 class TestLatentSpace(unittest.TestCase):
-    def test_sequence_keyframes(self):
-        num_frames = 20
-        batch_size = 3
-        keyframe_count = 3
-        dim_z = 128
-        dim_label = 1000
-        keyframes = [create_random_keyframe(dim_z,dim_label) for i in range(keyframe_count)]
-        z_seq, label_seq, truncation_seq = latent_space.sequence_keyframes(keyframes, num_frames, batch_size)
-        self.assertIs(len(z_seq), num_frames)
-        self.assertIs(len(label_seq), num_frames)
-        gan = biggan.BigGAN()
-        ims = gan.sample(z_seq, label_seq, truncation_seq, batch_size)
-        save_ims(ims)
+    def test_linear_interp_basic(self):
+        target_seq = np.asarray([
+            [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
+            [9, 8, 7, 6, 5, 4, 3, 2, 1, 0],
+            [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
+            ]).transpose()
+        points = np.asarray([target_seq[0], target_seq[-1]])
+        step_count = target_seq.shape[0]
+        actual_seq = latent_space.linear_interp(points, step_count)
+        compare_float_arrays_2d(target_seq, actual_seq)
+
+    def test_sequence_keyframes_linear_random_basic(self):
+        n_keyframes = 10
+        n_vector = 100
+        n_label = 1000
+        num_frames = 100
+        keyframes = np.asarray([create_random_keyframe(n_vector, n_label) for i in range(n_keyframes)])
+        z, label, trunc = latent_space.sequence_keyframes(
+                keyframes,
+                num_frames,
+                batch_size=1,
+                interp_method='linear')
+
+        assert (num_frames == z.shape[0]),\
+                'z sequence: target frame count: %s; actual shape: %s' % (num_frames, z.shape)
+        assert (num_frames == label.shape[0]),\
+                'label sequence: target frame count: %s; actual shape: %s' % (num_frames, label.shape)
+        assert (num_frames == trunc.shape[0]),\
+                'trunc sequence: target frame count: %s; actual shape: %s' % (num_frames, trunc.shape)
+
+    def test_sequence_keyframes_linear_random_batch(self):
+        n_keyframes = 10
+        n_vector = 100
+        n_label = 1000
+        num_frames = 100
+        batch_size = 7 # pick something that doesn't divide num_frames
+        batch_div = int(num_frames // batch_size)
+        batch_rem = 1 if int(num_frames % batch_size) > 0 else 0
+        batch_count = batch_div + batch_rem
+        keyframes = np.asarray([create_random_keyframe(n_vector, n_label) for i in range(n_keyframes)])
+        z, label, trunc = latent_space.sequence_keyframes(
+                keyframes,
+                num_frames,
+                batch_size=batch_size,
+                interp_method='linear')
+
+        assert (num_frames == z.shape[0]),\
+                'z sequence: target frame count: %s; actual shape: %s' % (num_frames, z.shape)
+        assert (num_frames == label.shape[0]),\
+                'label sequence: target frame count: %s; actual shape: %s' % (num_frames, label.shape)
+        assert (batch_count == trunc.shape[0]),\
+                'trunc sequence: target frame count: %s; actual shape: %s' % (batch_count, trunc.shape)
+
+    def test_sequence_keyframes_linear_random_batch_oob(self):
+        n_keyframes = 10
+        n_vector = 100
+        n_label = 1000
+        num_frames = 100
+        batch_size = 150 # pick something that doesn't divide num_frames
+        batch_div = int(num_frames // batch_size)
+        batch_rem = 1 if int(num_frames % batch_size) > 0 else 0
+        batch_count = batch_div + batch_rem
+        keyframes = np.asarray([create_random_keyframe(n_vector, n_label) for i in range(n_keyframes)])
+        z, label, trunc = latent_space.sequence_keyframes(
+                keyframes,
+                num_frames,
+                batch_size=batch_size,
+                interp_method='linear')
+
+        assert (num_frames == z.shape[0]),\
+                'z sequence: target frame count: %s; actual shape: %s' % (num_frames, z.shape)
+        assert (num_frames == label.shape[0]),\
+                'label sequence: target frame count: %s; actual shape: %s' % (num_frames, label.shape)
+        assert (batch_count == trunc.shape[0]),\
+                'trunc sequence: target frame count: %s; actual shape: %s' % (batch_count, trunc.shape)
+
+    def test_sequence_keyframes_cubic_random_basic(self):
+        n_keyframes = 10
+        n_vector = 100
+        n_label = 1000
+        num_frames = 100
+        keyframes = np.asarray([create_random_keyframe(n_vector, n_label) for i in range(n_keyframes)])
+        z, label, trunc = latent_space.sequence_keyframes(
+                keyframes,
+                num_frames,
+                batch_size=1,
+                interp_method='cubic')
+
+        assert (num_frames == z.shape[0]),\
+                'z sequence: target frame count: %s; actual shape: %s' % (num_frames, z.shape)
+        assert (num_frames == label.shape[0]),\
+                'label sequence: target frame count: %s; actual shape: %s' % (num_frames, label.shape)
+        assert (num_frames == trunc.shape[0]),\
+                'trunc sequence: target frame count: %s; actual shape: %s' % (num_frames, trunc.shape)
+
+    def test_sequence_keyframes_cubic_random_batch(self):
+        n_keyframes = 10
+        n_vector = 100
+        n_label = 1000
+        num_frames = 100
+        batch_size = 7 # pick something that doesn't divide num_frames
+        batch_div = int(num_frames // batch_size)
+        batch_rem = 1 if int(num_frames % batch_size) > 0 else 0
+        batch_count = batch_div + batch_rem
+        keyframes = np.asarray([create_random_keyframe(n_vector, n_label) for i in range(n_keyframes)])
+        z, label, trunc = latent_space.sequence_keyframes(
+                keyframes,
+                num_frames,
+                batch_size=batch_size,
+                interp_method='cubic')
+
+        assert (num_frames == z.shape[0]),\
+                'z sequence: target frame count: %s; actual shape: %s' % (num_frames, z.shape)
+        assert (num_frames == label.shape[0]),\
+                'label sequence: target frame count: %s; actual shape: %s' % (num_frames, label.shape)
+        assert (batch_count == trunc.shape[0]),\
+                'trunc sequence: target frame count: %s; actual shape: %s' % (batch_count, trunc.shape)
 
+    def test_sequence_keyframes_cubic_random_batch_oob(self):
+        n_keyframes = 10
+        n_vector = 100
+        n_label = 1000
+        num_frames = 100
+        batch_size = 150 # pick something that doesn't divide num_frames
+        batch_div = int(num_frames // batch_size)
+        batch_rem = 1 if int(num_frames % batch_size) > 0 else 0
+        batch_count = batch_div + batch_rem
+        keyframes = np.asarray([create_random_keyframe(n_vector, n_label) for i in range(n_keyframes)])
+        z, label, trunc = latent_space.sequence_keyframes(
+                keyframes,
+                num_frames,
+                batch_size=batch_size,
+                interp_method='cubic')
 
 
 if __name__ == '__main__':