def make_grid(samples, show_num=64):
  """Tile images to an image grid."""
  batch_size, height, width, c = samples.shape
  if batch_size < show_num:
    logging.info('show_num is cut by the small batch size to %d', batch_size)
    show_num = batch_size
  h_num = int(math.sqrt(show_num))
  w_num = int(show_num / h_num)
  grid_num = h_num * w_num

  samples = samples[0:grid_num]
  samples = samples.reshape(h_num, w_num, height, width, c)
  samples = samples.swapaxes(1, 2)
  samples = samples.reshape(height * h_num, width * w_num, c)
  return samples
Tile images to an image grid.
videogvt/train_lib/train_utils.py