visualize.py

# ------------------------------------------------------------------------------
# Util functions to visualize images.
# ------------------------------------------------------------------------------

import numpy as np
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, [a, b, h, w, c])
  x = np.transpose(x, [0, 2, 1, 3, 4])
  x = np.reshape(x, [a * h, b * w, 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)