refactor and add hdf5 support

3 files changed, 107 insertions, 78 deletions

diff --git a/cli/app/search/image.py b/cli/app/search/image.py
new file mode 100644
index 0000000..f800a33
--- /dev/null
+++ b/cli/app/search/image.py
@@ -0,0 +1,87 @@
+import cv2 as cv
+import numpy as np
+
+def image_to_uint8(x):
+  """Converts [-1, 1] float array to [0, 255] uint8."""
+  x = np.asarray(x)
+  x = (256. / 2.) * (x + 1.)
+  x = np.clip(x, 0, 255)
+  x = x.astype(np.uint8)
+  return x
+
+def imconvert_uint8(im):
+  im = np.clip(((im + 1) / 2.0) * 256, 0, 255)
+  im = np.uint8(im)
+  return im
+
+def imconvert_float32(im):
+  im = np.float32(im)
+  im = (im / 256) * 2.0 - 1
+  return im
+
+def imread(filename):
+  img = cv.imread(filename, cv.IMREAD_UNCHANGED)
+  if img is not None:
+    if len(img.shape) > 2:
+      img = img[...,::-1]
+  return img
+  
+def imwrite(filename, img):
+  if img is not None:
+    if len(img.shape) > 2:
+      img = img[...,::-1]
+  return cv.imwrite(filename, img)
+
+def imgrid(imarray, cols=5, pad=1):
+  if imarray.dtype != np.uint8:
+    raise ValueError('imgrid input imarray must be uint8')
+  pad = int(pad)
+  assert pad >= 0
+  cols = int(cols)
+  assert cols >= 1
+  N, H, W, C = imarray.shape
+  rows = int(np.ceil(N / float(cols)))
+  batch_pad = rows * cols - N
+  assert batch_pad >= 0
+  post_pad = [batch_pad, pad, pad, 0]
+  pad_arg = [[0, p] for p in post_pad]
+  imarray = np.pad(imarray, pad_arg, 'constant', constant_values=255)
+  H += pad
+  W += pad
+  grid = (imarray
+          .reshape(rows, cols, H, W, C)
+          .transpose(0, 2, 1, 3, 4)
+          .reshape(rows*H, cols*W, C))
+  if pad:
+    grid = grid[:-pad, :-pad]
+  return grid
+
+def resize_and_crop_image(target_im, opt_dims):
+  w = target_im.shape[1]
+  h = target_im.shape[0]
+  if w <= h:
+    scale = opt_dims / w
+  else:
+    scale = opt_dims / h
+  #print("{} {}".format(w, h))
+  target_im = cv.resize(target_im,(0,0), fx=scale, fy=scale)
+
+  w = target_im.shape[1]
+  h = target_im.shape[0]
+
+  x0 = 0
+  x1 = opt_dims
+  y0 = 0
+  y1 = opt_dims
+  if w > opt_dims:
+    x0 += int((w - opt_dims) / 2)
+    x1 += x0
+  if h > opt_dims:
+    y0 += int((h - opt_dims) / 2)
+    y1 += y0
+
+  phi_target = imconvert_float32(target_im)
+  phi_target = phi_target[y0:y1,x0:x1]
+  if phi_target.shape[2] == 4:
+    phi_target = phi_target[:,:,1:4]
+  return phi_target
diff --git a/cli/app/search/util.py b/cli/app/search/util.py
deleted file mode 100644
index a4cdfd9..0000000
--- a/cli/app/search/util.py
+++ /dev/null
@@ -1,78 +0,0 @@
-
-def truncated_z_sample(batch_size, truncation=1., seed=None):
-  state = None if seed is None else np.random.RandomState(seed)
-  values = truncnorm.rvs(-2, 2, size=(batch_size, dim_z), random_state=state)
-  return truncation * values
-
-def one_hot(index, vocab_size=vocab_size):
-  index = np.asarray(index)
-  if len(index.shape) == 0:
-    index = np.asarray([index])
-  assert len(index.shape) == 1
-  num = index.shape[0]
-  output = np.zeros((num, vocab_size), dtype=np.float32)
-  output[np.arange(num), index] = 1
-  return output
-
-def imconvert_uint8(im):
-  im = np.clip(((im + 1) / 2.0) * 256, 0, 255)
-  im = np.uint8(im)
-  return im
-
-def imconvert_float32(im):
-  im = np.float32(im)
-  im = (im / 256) * 2.0 - 1
-  return im
-
-def imread(filename):
-  img = cv2.imread(filename, cv2.IMREAD_UNCHANGED)
-  if img is not None:
-    if len(img.shape) > 2:
-      img = img[...,::-1]
-  return img
-  
-def imwrite(filename, img):
-  if img is not None:
-    if len(img.shape) > 2:
-      img = img[...,::-1]
-  return cv2.imwrite(filename, img)
-
-def imgrid(imarray, cols=5, pad=1):
-  if imarray.dtype != np.uint8:
-    raise ValueError('imgrid input imarray must be uint8')
-  pad = int(pad)
-  assert pad >= 0
-  cols = int(cols)
-  assert cols >= 1
-  N, H, W, C = imarray.shape
-  rows = int(np.ceil(N / float(cols)))
-  batch_pad = rows * cols - N
-  assert batch_pad >= 0
-  post_pad = [batch_pad, pad, pad, 0]
-  pad_arg = [[0, p] for p in post_pad]
-  imarray = np.pad(imarray, pad_arg, 'constant', constant_values=255)
-  H += pad
-  W += pad
-  grid = (imarray
-          .reshape(rows, cols, H, W, C)
-          .transpose(0, 2, 1, 3, 4)
-          .reshape(rows*H, cols*W, C))
-  if pad:
-    grid = grid[:-pad, :-pad]
-  return grid
-
-def imshow(a, format='png', jpeg_fallback=True):
-  a = np.asarray(a, dtype=np.uint8)
-  str_file = cStringIO.StringIO()
-  PIL.Image.fromarray(a).save(str_file, format)
-  im_data = str_file.getvalue()
-  try:
-    disp = IPython.display.display(IPython.display.Image(im_data))
-  except IOError:
-    if jpeg_fallback and format != 'jpeg':
-      print ('Warning: image was too large to display in format "{}"; '
-             'trying jpeg instead.').format(format)
-      return imshow(a, format='jpeg')
-    else:
-      raise
-  return disp
diff --git a/cli/app/search/vector.py b/cli/app/search/vector.py
new file mode 100644
index 0000000..89cd949
--- /dev/null
+++ b/cli/app/search/vector.py
@@ -0,0 +1,20 @@
+import random
+import numpy as np
+from scipy.stats import truncnorm
+
+def truncated_z_sample(batch_size, z_dim, truncation):
+  values = truncnorm.rvs(-2, 2, size=(batch_size, z_dim))
+  return truncation * values
+
+def truncated_z_single(z_dim, truncation):
+  values = truncnorm.rvs(-2, 2, size=(1, z_dim))
+  return truncation * values
+
+def create_labels(batch_size, vocab_size, num_classes):
+  label = np.zeros((batch_size, vocab_size))
+  for i in range(batch_size):
+    for _ in range(random.randint(1, num_classes)):
+      j = random.randint(0, vocab_size-1)
+      label[i, j] = random.random()
+    label[i] /= label[i].sum()
+  return label