# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Image util ops."""

import tensorflow as tf


def get_ndims(image):
    return image.get_shape().ndims or tf.rank(image)


def to_4D_image(image):
    """Convert 2/3/4D image to 4D image.

    Args:
      image: 2/3/4D tensor.

    Returns:
      4D tensor with the same type.
    """
    with tf.control_dependencies(
        [
            tf.debugging.assert_rank_in(
                image, [2, 3, 4], message="`image` must be 2/3/4D tensor"
            )
        ]
    ):
        ndims = image.get_shape().ndims
        if ndims is None:
            return _dynamic_to_4D_image(image)
        elif ndims == 2:
            return image[None, :, :, None]
        elif ndims == 3:
            return image[None, :, :, :]
        else:
            return image


def _dynamic_to_4D_image(image):
    shape = tf.shape(image)
    original_rank = tf.rank(image)
    # 4D image => [N, H, W, C] or [N, C, H, W]
    # 3D image => [1, H, W, C] or [1, C, H, W]
    # 2D image => [1, H, W, 1]
    left_pad = tf.cast(tf.less_equal(original_rank, 3), dtype=tf.int32)
    right_pad = tf.cast(tf.equal(original_rank, 2), dtype=tf.int32)
    new_shape = tf.concat(
        [
            tf.ones(shape=left_pad, dtype=tf.int32),
            shape,
            tf.ones(shape=right_pad, dtype=tf.int32),
        ],
        axis=0,
    )
    return tf.reshape(image, new_shape)


def from_4D_image(image, ndims):
    """Convert back to an image with `ndims` rank.

    Args:
      image: 4D tensor.
      ndims: The original rank of the image.

    Returns:
      `ndims`-D tensor with the same type.
    """
    with tf.control_dependencies(
        [tf.debugging.assert_rank(image, 4, message="`image` must be 4D tensor")]
    ):
        if isinstance(ndims, tf.Tensor):
            return _dynamic_from_4D_image(image, ndims)
        elif ndims == 2:
            return tf.squeeze(image, [0, 3])
        elif ndims == 3:
            return tf.squeeze(image, [0])
        else:
            return image


def _dynamic_from_4D_image(image, original_rank):
    shape = tf.shape(image)
    # 4D image <= [N, H, W, C] or [N, C, H, W]
    # 3D image <= [1, H, W, C] or [1, C, H, W]
    # 2D image <= [1, H, W, 1]
    begin = tf.cast(tf.less_equal(original_rank, 3), dtype=tf.int32)
    end = 4 - tf.cast(tf.equal(original_rank, 2), dtype=tf.int32)
    new_shape = shape[begin:end]
    return tf.reshape(image, new_shape)

[((1499, 1514), 'tensorflow.shape', 'tf.shape', (['image'], {}), '(image)\n', (1507, 1514), True, 'import tensorflow as tf\n'), ((1535, 1549), 'tensorflow.rank', 'tf.rank', (['image'], {}), '(image)\n', (1542, 1549), True, 'import tensorflow as tf\n'), ((2022, 2050), 'tensorflow.reshape', 'tf.reshape', (['image', 'new_shape'], {}), '(image, new_shape)\n', (2032, 2050), True, 'import tensorflow as tf\n'), ((2742, 2757), 'tensorflow.shape', 'tf.shape', (['image'], {}), '(image)\n', (2750, 2757), True, 'import tensorflow as tf\n'), ((3062, 3090), 'tensorflow.reshape', 'tf.reshape', (['image', 'new_shape'], {}), '(image, new_shape)\n', (3072, 3090), True, 'import tensorflow as tf\n'), ((798, 812), 'tensorflow.rank', 'tf.rank', (['image'], {}), '(image)\n', (805, 812), True, 'import tensorflow as tf\n'), ((1698, 1729), 'tensorflow.less_equal', 'tf.less_equal', (['original_rank', '(3)'], {}), '(original_rank, 3)\n', (1711, 1729), True, 'import tensorflow as tf\n'), ((1771, 1797), 'tensorflow.equal', 'tf.equal', (['original_rank', '(2)'], {}), '(original_rank, 2)\n', (1779, 1797), True, 'import tensorflow as tf\n'), ((2903, 2934), 'tensorflow.less_equal', 'tf.less_equal', (['original_rank', '(3)'], {}), '(original_rank, 3)\n', (2916, 2934), True, 'import tensorflow as tf\n'), ((1864, 1903), 'tensorflow.ones', 'tf.ones', ([], {'shape': 'left_pad', 'dtype': 'tf.int32'}), '(shape=left_pad, dtype=tf.int32)\n', (1871, 1903), True, 'import tensorflow as tf\n'), ((1936, 1976), 'tensorflow.ones', 'tf.ones', ([], {'shape': 'right_pad', 'dtype': 'tf.int32'}), '(shape=right_pad, dtype=tf.int32)\n', (1943, 1976), True, 'import tensorflow as tf\n'), ((2974, 3000), 'tensorflow.equal', 'tf.equal', (['original_rank', '(2)'], {}), '(original_rank, 2)\n', (2982, 3000), True, 'import tensorflow as tf\n'), ((1033, 1124), 'tensorflow.debugging.assert_rank_in', 'tf.debugging.assert_rank_in', (['image', '[2, 3, 4]'], {'message': '"""`image` must be 2/3/4D tensor"""'}), "(image, [2, 3, 4], message=\n    '`image` must be 2/3/4D tensor')\n", (1060, 1124), True, 'import tensorflow as tf\n'), ((2325, 2396), 'tensorflow.debugging.assert_rank', 'tf.debugging.assert_rank', (['image', '(4)'], {'message': '"""`image` must be 4D tensor"""'}), "(image, 4, message='`image` must be 4D tensor')\n", (2349, 2396), True, 'import tensorflow as tf\n'), ((2546, 2571), 'tensorflow.squeeze', 'tf.squeeze', (['image', '[0, 3]'], {}), '(image, [0, 3])\n', (2556, 2571), True, 'import tensorflow as tf\n'), ((2616, 2638), 'tensorflow.squeeze', 'tf.squeeze', (['image', '[0]'], {}), '(image, [0])\n', (2626, 2638), True, 'import tensorflow as tf\n')]