import unittest import numpy as np from scipy.stats import truncnorm import context from gantools import latent_space import PIL.Image import math def create_random_keyframe(n_vector, n_label): truncation = (0.9 - 0.1)*np.random.random() + 0.1 random_state = np.random.RandomState() vectors = truncnorm.rvs(-2, 2, size=(n_vector,), random_state=random_state) keyframe = { 'vector': vectors.tolist(), 'label': latent_space.one_hot(np.random.randint(0, n_label), n_label), 'truncation': truncation, } return keyframe #### TMP def save_image(arr, fp): image = PIL.Image.fromarray(arr) image.save(fp, format='JPEG', quality=90) def save_ims(ims): i = 0 for im in ims: path = './GAN_'+str(i).zfill(3)+'.jpeg' save_image(im, path) 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_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__': unittest.main()