def testDepthInterleaved(self):
   x_np = [[[[1, 10], [2, 20]],
            [[3, 30], [4, 40]]]]
   with self.test_session(use_gpu=False):
     block_size = 2
     out_tf = tf.space_to_depth(x_np, block_size)
     self.assertAllEqual(out_tf.eval(), [[[[1, 10, 2, 20, 3, 30, 4, 40]]]])

 def testBlockSizeNotDivisibleWidth(self):
   # The block size divides width but not height.
   x_np = [[[[1], [2], [3]],
            [[3], [4], [7]]]]
   block_size = 3
   with self.assertRaises(IndexError):
     _ = tf.space_to_depth(x_np, block_size)

 def testBlockSizeNotDivisibleBoth(self):
   # The block size does not divide neither width or height.
   x_np = [[[[1], [2]],
            [[3], [4]]]]
   block_size = 3
   with self.assertRaises(IndexError):
     _ = tf.space_to_depth(x_np, block_size)

 def testBasic(self):
   x_np = [[[[1], [2]],
            [[3], [4]]]]
   with self.test_session(use_gpu=False):
     block_size = 2
     out_tf = tf.space_to_depth(x_np, block_size)
     self.assertAllEqual(out_tf.eval(), [[[[1, 2, 3, 4]]]])

 def testInputWrongDimMissingBatch(self):
   # The input is missing the first dimension ("batch")
   x_np = [[[1], [2]],
           [[3], [4]]]
   block_size = 2
   with self.assertRaises(ValueError):
     _ = tf.space_to_depth(x_np, block_size)

 def testBlockSize0(self):
   # The block size is 0.
   x_np = [[[[1], [2]],
            [[3], [4]]]]
   block_size = 0
   with self.assertRaises(ValueError):
     out_tf = tf.space_to_depth(x_np, block_size)
     out_tf.eval()

 def testBlockSizeOne(self):
   # The block size is 1. The block size needs to be > 1.
   x_np = [[[[1], [2]],
            [[3], [4]]]]
   block_size = 1
   with self.assertRaises(ValueError):
     out_tf = tf.space_to_depth(x_np, block_size)
     out_tf.eval()

 def testBlockSizeLarger(self):
   # The block size is too large for this input.
   x_np = [[[[1], [2]],
            [[3], [4]]]]
   block_size = 10
   with self.assertRaises(IndexError):
     out_tf = tf.space_to_depth(x_np, block_size)
     out_tf.eval()

 def testBlockSizeNotDivisibleHeight(self):
   # The block size divides height but not width.
   x_np = [[[[1], [2]],
            [[3], [4]],
            [[5], [6]]]]
   block_size = 3
   with self.assertRaises(IndexError):
     _ = tf.space_to_depth(x_np, block_size)

 def testSpaceToDepthTranspose(self):
     x = np.arange(5 * 10 * 16 * 7, dtype=np.float32).reshape([5, 10, 16, 7])
     block_size = 2
     paddings = np.zeros((2, 2), dtype=np.int32)
     y1 = tf.space_to_batch(x, paddings, block_size=block_size)
     y2 = tf.transpose(tf.space_to_depth(tf.transpose(x, [3, 1, 2, 0]), block_size=block_size), [3, 1, 2, 0])
     with self.test_session():
         self.assertAllEqual(y1.eval(), y2.eval())

 def testDepthInterleavedDepth3(self):
   x_np = [[[[1, 2, 3], [4, 5, 6]],
            [[7, 8, 9], [10, 11, 12]]]]
   with self.test_session(use_gpu=False):
     block_size = 2
     out_tf = tf.space_to_depth(x_np, block_size)
     self.assertAllEqual(out_tf.eval(),
                         [[[[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]]]])

 def testInputWrongDimMissingDepth(self):
   # The input is missing the last dimension ("depth")
   x_np = [[[1, 2],
            [3, 4]]]
   block_size = 2
   with self.assertRaises(ValueError):
     out_tf = tf.space_to_depth(x_np, block_size)
     out_tf.eval()

 def testBlockSizeNotDivisibleDepth(self):
   # The depth is not divisible by the square of the block size.
   x_np = [[[[1, 1, 1, 1],
             [2, 2, 2, 2]],
            [[3, 3, 3, 3],
             [4, 4, 4, 4]]]]
   block_size = 3
   with self.assertRaises(IndexError):
     _ = tf.space_to_depth(x_np, block_size)

  def testLargerInput4x4(self):
    x_np = [[[[1], [2], [5], [6]],
             [[3], [4], [7], [8]],
             [[9], [10], [13], [14]],
             [[11], [12], [15], [16]]]]

    with self.test_session(use_gpu=False):
      block_size = 4
      out_tf = tf.space_to_depth(x_np, block_size)
      self.assertAllEqual(
          out_tf.eval(),
          [[[[1, 2, 5, 6, 3, 4, 7, 8, 9, 10, 13, 14, 11, 12, 15, 16]]]])

 def testDepthInterleavedLarge(self):
   x_np = [[[[1, 10], [2, 20], [5, 50], [6, 60]],
            [[3, 30], [4, 40], [7, 70], [8, 80]],
            [[9, 90], [10, 100], [13, 130], [14, 140]],
            [[11, 110], [12, 120], [15, 150], [16, 160]]]]
   with self.test_session(use_gpu=False):
     block_size = 2
     out_tf = tf.space_to_depth(x_np, block_size)
     self.assertAllEqual(out_tf.eval(),
                         [[[[1, 10, 2, 20, 3, 30, 4, 40],
                            [5, 50, 6, 60, 7, 70, 8, 80]],
                           [[9, 90, 10, 100, 11, 110, 12, 120],
                            [13, 130, 14, 140, 15, 150, 16, 160]]]])

 def testNonSquare(self):
   x_np = [[[[1, 10], [2, 20]],
            [[3, 30], [4, 40]],
            [[5, 50], [6, 60]],
            [[7, 70], [8, 80]],
            [[9, 90], [10, 100]],
            [[11, 110], [12, 120]]]]
   with self.test_session(use_gpu=False):
     block_size = 2
     out_tf = tf.space_to_depth(x_np, block_size)
     self.assertAllEqual(out_tf.eval(),
                         [[[[1, 10, 2, 20, 3, 30, 4, 40]],
                           [[5, 50, 6, 60, 7, 70, 8, 80]],
                           [[9, 90, 10, 100, 11, 110, 12, 120]]]])

  def _checkGrad(self, x, block_size):
    assert 4 == x.ndim
    with self.test_session():
      tf_x = tf.convert_to_tensor(x)
      tf_y = tf.space_to_depth(tf_x, block_size)
      epsilon = 1e-2
      ((x_jacob_t, x_jacob_n)) = tf.test.compute_gradient(
          tf_x,
          x.shape,
          tf_y,
          tf_y.get_shape().as_list(),
          x_init_value=x,
          delta=epsilon)

    self.assertAllClose(x_jacob_t, x_jacob_n, rtol=1e-2, atol=epsilon)

 def __call__(self, shape, dtype='float32'):  # tf needs partition_info=None
     shape = list(shape)
     if self.scale == 1:
         return self.initializer(shape)
     new_shape = shape[:3] + [shape[3] // (self.scale ** 2)]
     if type(self.initializer) is dict:
         self.initializer = initializers.deserialize(self.initializer)
     var_x = self.initializer(new_shape, dtype)
     var_x = tf.transpose(var_x, perm=[2, 0, 1, 3])
     var_x = tf.image.resize_nearest_neighbor(
                      var_x,
                      size=(shape[0] * self.scale, shape[1] * self.scale),
                      align_corners=True)
     var_x = tf.space_to_depth(var_x, block_size=self.scale, data_format='NHWC')
     var_x = tf.transpose(var_x, perm=[1, 2, 0, 3])
     return var_x

def icnr_keras(shape, dtype=None):
    """
    Custom initializer for subpix upscaling
    From https://github.com/kostyaev/ICNR
    Note: upscale factor is fixed to 2, and the base initializer is fixed to random normal.
    """
    # TODO Roll this into ICNR_init when porting GAN 2.2
    shape = list(shape)
    scale = 2
    initializer = tf.keras.initializers.RandomNormal(0, 0.02)

    new_shape = shape[:3] + [int(shape[3] / (scale ** 2))]
    var_x = initializer(new_shape, dtype)
    var_x = tf.transpose(var_x, perm=[2, 0, 1, 3])
    var_x = tf.image.resize_nearest_neighbor(var_x, size=(shape[0] * scale, shape[1] * scale))
    var_x = tf.space_to_depth(var_x, block_size=scale)
    var_x = tf.transpose(var_x, perm=[1, 2, 0, 3])
    return var_x

def read_and_batchify_image(image_path, shape, image_type="jpg"):
    """Return the original image as read from image_path and the image splitted as a batch tensor.
    Args:
        image_path: image path
        shape: batch shape, like: [no_patches_per_side**2, patch_side, patch_side, 3]
        image_type: image type
    Returns:
        original_image, patches
        where original image is a tensor in the format [widht, height 3]
        and patches is a tensor of processed images, ready to be classified, with size
        [batch_size, w, h, 3]"""

    original_image = read_image(image_path, 3, image_type)

    # extract values from shape
    patch_side = shape[1]
    no_patches_per_side = int(math.sqrt(shape[0]))
    resized_input_side = patch_side * no_patches_per_side

    resized_image = resize_bl(original_image, resized_input_side)

    resized_image = tf.expand_dims(resized_image, 0)
    patches = tf.space_to_depth(resized_image, patch_side)
    print(patches)
    patches = tf.squeeze(patches, [0])  #4,4,192*192*3
    print(patches)
    patches = tf.reshape(patches,
                         [no_patches_per_side**2, patch_side, patch_side, 3])
    print(patches)
    patches_a = tf.split(0, no_patches_per_side**2, patches)
    print(patches_a)
    normalized_patches = []
    for patch in patches_a:
        patch_as_input_image = zm_mp(
            tf.reshape(tf.squeeze(patch, [0]), [patch_side, patch_side, 3]))
        print(patch_as_input_image)
        normalized_patches.append(patch_as_input_image)

    # the last patch is not a "patch" but the whole image resized to patch_side² x 3
    # to give a glance to the whole image, in parallel with the patch analysis
    normalized_patches.append(zm_mp(resize_bl(original_image, patch_side)))
    batch_of_patches = tf.pack(normalized_patches)
    return tf.image.convert_image_dtype(original_image,
                                        tf.uint8), batch_of_patches