1 files changed, 88 insertions, 0 deletions
diff --git a/inversion/visualize.py b/inversion/visualize.py
new file mode 100644
index 0000000..07aea2d
--- /dev/null
+++ b/inversion/visualize.py
@@ -0,0 +1,88 @@
+# ------------------------------------------------------------------------------
+# Util functions to visualize images.
+# ------------------------------------------------------------------------------
+
+import numpy as np
+import cv2 as cv
+from PIL import Image
+
+def split(x):
+  assert type(x) == int
+  t = int(np.floor(np.sqrt(x)))
+  for a in range(t, 0, -1):
+    if x % a == 0:
+      return a, x / a
+
+def grid_transform(x):
+  n, c, h, w = x.shape
+  a, b = split(n)
+  x = np.transpose(x, [0, 2, 3, 1])
+  x = np.reshape(x, [int(a), int(b), int(h), int(w), int(c)])
+  x = np.transpose(x, [0, 2, 1, 3, 4])
+  x = np.reshape(x, [int(a * h), int(b * w), int(c)])
+  if x.shape[2] == 1:
+    x = np.squeeze(x, axis=2)
+  return x
+
+def seq_transform(x):
+  n, c, h, w = x.shape
+  x = np.transpose(x, [2, 0, 3, 1])
+  x = np.reshape(x, [h, n * w, c])
+  return x
+
+# Converts image pixels from range [-1, 1] to [0, 255].
+def data2img(data):
+  rescaled = np.divide(data + 1.0, 2.0) * 255.
+  rescaled = np.clip(rescaled, 0, 255)
+  return np.rint(rescaled).astype('uint8')
+
+def interleave(a, b):
+  res = np.empty([a.shape[0] + b.shape[0]] + list(a.shape[1:]), dtype=a.dtype)
+  res[0::2] = a
+  res[1::2] = b
+  return res
+
+def save_image(filepath, img):
+  pilimg = Image.fromarray(img)
+  pilimg.save(filepath)
+
+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 imconvert_float32(im):
+  im = np.float32(im)
+  im = (im / 256) * 2.0 - 1
+  return im
+
+def load_image(opt_fp_in, opt_dims=128):
+  target_im = imread(opt_fp_in)
+  w = target_im.shape[1]
+  h = target_im.shape[0]
+  if w <= h:
+    scale = opt_dims / w
+  else:
+    scale = opt_dims / 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]
+  phi_target = np.expand_dims(phi_target, 0)
+  return phi_target