sampling

1 files changed, 31 insertions, 5 deletions

diff --git a/inversion/live.py b/inversion/live.py
index ad4b80f..491d4a8 100644
--- a/inversion/live.py
+++ b/inversion/live.py
@@ -8,6 +8,8 @@ import tensorflow as tf
 import tensorflow_probability as tfp
 import tensorflow_hub as hub
 import time
+import random
+from scipy.stats import truncnorm
 from PIL import Image
 import visualize as vs
 tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
@@ -16,6 +18,8 @@ from rpc import CortexRPC
 
 from params import Params
 
+FPS = 25
+
 params = Params('params_dense-512.json')
 
 # --------------------------
@@ -46,8 +50,7 @@ N_CLASS = input_info['y'].get_shape().as_list()[1]
 # Initializers
 # --------------------------
 
-def label_sampler(shape=[BATCH_SIZE, N_CLASS]):
-  num_classes = 2
+def label_sampler(num_classes=2, shape=(BATCH_SIZE, N_CLASS,)):
   label = np.zeros(shape)
   for i in range(shape[0]):
     for _ in range(random.randint(1, shape[1])):
@@ -56,6 +59,13 @@ def label_sampler(shape=[BATCH_SIZE, N_CLASS]):
     label[i] /= label[i].sum()
   return label
 
+def truncated_z_sample(shape=(BATCH_SIZE, Z_DIM,), truncation=1.0):
+  values = truncnorm.rvs(-2, 2, size=shape)
+  return truncation * values
+
+def normal_z_sample(shape=(BATCH_SIZE, Z_DIM,)):
+  return np.random.normal(size=shape)
+
 # --------------------------
 # More complex ops
 # --------------------------
@@ -63,21 +73,29 @@ def label_sampler(shape=[BATCH_SIZE, N_CLASS]):
 def sin(opts, key, shape):
   noise = lerp(opts, key + '_noise', shape)
   scale = InterpolatorParam(name=key + '_scale')
+  speed = InterpolatorParam(name=key + '_speed')
   time = opts['time'].variable
   out = tf.sin(time + noise) * scale.variable
   opts[key + '_scale'] = scale
+  opts[key + '_speed'] = speed
   return out
 
 def lerp(opts, key, shape):
   a = InterpolatorParam(name=key + '_a', shape=shape)
   b = InterpolatorParam(name=key + '_b', shape=shape)
   n = InterpolatorParam(name=key + '_n')
+  speed = InterpolatorParam(name=key + '_speed')
   out = a.variable * (1 - n.variable) + b.variable * n.variable
   opts[key + '_a'] = a
   opts[key + '_b'] = b
   opts[key + '_n'] = n
+  opts[key + '_speed'] = speed
   return out
 
+# --------------------------
+# Placeholder params
+# --------------------------
+
 class InterpolatorParam:
   def __init__(self, name, dtype=tf.float32, shape=(), value=None, type="noise"):
     self.scalar = shape == ()
@@ -91,11 +109,15 @@ class InterpolatorParam:
 
   def randomize(self, num_classes):
     if self.type == 'noise':
-      val = np.random.normal(size=self.shape)
+      val = truncated_z_sample(shape=self.shape)
     elif self.type == 'label':
       val = label_sampler(shape=self.shape)
     self.assign(val)
 
+# --------------------------
+# Interpolator graph
+# --------------------------
+
 class Interpolator:
   def __init__(self):
     self.opts = {
@@ -105,7 +127,7 @@ class Interpolator:
 
   def build(self):
     lerp_z = lerp(self.opts, 'latent', [BATCH_SIZE, Z_DIM])
-    sin_z = sin(self.opts, 'sin_z', [BATCH_SIZE, Z_DIM])
+    sin_z = sin(self.opts, 'orbit', [BATCH_SIZE, Z_DIM])
     lerp_label = lerp(self.opts, 'label', [BATCH_SIZE, N_CLASS])
 
     gen_in = {}
@@ -141,9 +163,12 @@ class Interpolator:
     # do things like create a new B and interpolate to it
     pass
 
+# --------------------------
+# RPC Listener
+# --------------------------
+
 class Listener:
   def __init__(self):
-    self.assign_ops = {}
     self.interpolator = Interpolator()
     self.interpolator.build()
 
@@ -167,6 +192,7 @@ class Listener:
     print("Ready!")
     self.rpc_client = rpc_client
     self.rpc_client.send_status('processing', True)
+    dt = 1 / FPS
     for i in range(99999):
       print("Step {}".format(i))
       gen_time = time.time()