utils.py

from PIL import Image
import tensorflow as tf
import glob
import json
import numpy as np
import os
import pdb

def batch_reader_list(img_names, index, label_map, img_shape, batch_size=1):    
    """ Gets the names of the files and ground truth for images and converts them
        to a tf object.
    """
    img_list = []
    ground_truth_all_classes = []
    ground_truth_all_bboxes = []
    for batch_index in range(0, batch_size):
        img = np.asarray(Image.open(img_names[index+batch_index]))
        img_list.append(img)
        ground_truth_name = '{}.{}'.format(os.path.splitext(img_names[index+batch_index])[0], 'json')
        with open(ground_truth_name) as f:
            ground_truth = json.load(f)

        ground_truth_class_tensor = np.zeros((len(ground_truth)), np.int64)
        ground_truth_bbox_tensor = np.zeros((len(ground_truth), 4), np.float32)
        for ind, gt_inn in enumerate(ground_truth):
            ground_truth_class_tensor[ind] = label_map.index(gt_inn[1])
            ground_truth_bbox_tensor[ind, :] = [gt_inn[0][1], gt_inn[0][0], gt_inn[0][3], gt_inn[0][2]]
        ground_truth_all_bboxes.append(ground_truth_bbox_tensor)
        ground_truth_all_classes.append(ground_truth_class_tensor)

    return img_list, ground_truth_all_bboxes, ground_truth_all_classes

def batch_reader(img_names, index, label_map, img_shape, batch_size=1):    
    """ Gets the names of the files and ground truth for images and converts them
        to a tf object.
    """
    img_tensor = np.zeros((batch_size, img_shape[0], img_shape[1], 3), dtype=np.int32)
    ground_truth_all_classes = []
    ground_truth_all_bboxes = []
    for batch_index in range(0, batch_size):
        img = np.asarray(Image.open(img_names[index+batch_index]))
        img_tensor[batch_index, :, :, :] = img
        ground_truth_name = '{}.{}'.format(os.path.splitext(img_names[index+batch_index])[0], 'json')
        with open(ground_truth_name) as f:
            ground_truth = json.load(f)

        ground_truth_class_tensor = np.zeros((len(ground_truth)), np.int64)
        ground_truth_bbox_tensor = np.zeros((len(ground_truth), 4), np.float32)
        for ind, gt_inn in enumerate(ground_truth):
            ground_truth_class_tensor[ind] = label_map.index(gt_inn[1])
            ground_truth_bbox_tensor[ind, :] = [gt_inn[0][1], gt_inn[0][0], gt_inn[0][3], gt_inn[0][2]]
        ground_truth_all_bboxes.append(ground_truth_bbox_tensor)
        ground_truth_all_classes.append(ground_truth_class_tensor)

    return img_tensor, ground_truth_all_bboxes, ground_truth_all_classes

def parse_function(filename):
    """ Reads an image from a file, decodes it into a dense tensor, and resizes it
    to a fixed shape """
    image_string = tf.read_file(filename)
    image_decoded = tf.image.decode_jpeg(image_string)
    image_scaled = tf.image.per_image_standardization(image_decoded)
    return image_scaled

def create_tf_dataset(file_names, gt_bboxes, gt_classes, buffer_size, number_iterations_dataset, batch_size=1):
    """ This function returns a pointer to iterate over the batches of data """
    train_dataset = tf.data.Dataset.from_tensor_slices((file_names, gt_bboxes, gt_classes))
    train_dataset = train_dataset.map(parse_function)
    train_dataset = train_dataset.repeat(number_iterations_dataset)
    train_dataset = train_dataset.shuffle(buffer_size=buffer_size)
    train_batched_dataset = train_dataset.batch(batch_size)
    train_iterator = train_batched_dataset.make_initializable_iterator()

    return train_iterator.get_next(), train_iterator

def ssd_bboxes_encode_layer(labels,
                           bboxes,
                           anchors_layer,
                           num_classes,
                           prior_scaling=[0.1, 0.1, 0.2, 0.2],
                           dtype=tf.float32):
    """Encode groundtruth labels and bounding boxes using SSD anchors from
    one layer.
    Arguments:
      labels: 1D Tensor(int64) containing groundtruth labels;
      bboxes: Nx4 Tensor(float) with bboxes relative coordinates;
      anchors_layer: Numpy array with layer anchors;
      matching_threshold: Threshold for positive match with groundtruth bboxes;
      prior_scaling: Scaling of encoded coordinates.
    Return:
      (target_labels, target_localizations, target_scores): Target Tensors.
    """
    # Anchors coordinates and volume.
    yref, xref, href, wref = anchors_layer
    ymin = yref - href / 2.
    xmin = xref - wref / 2.
    ymax = yref + href / 2.
    xmax = xref + wref / 2.
    vol_anchors = (xmax - xmin) * (ymax - ymin)

    # Initialize tensors...
    shape = (yref.shape[0], yref.shape[1], href.size)
    feat_labels = tf.zeros(shape, dtype=tf.int64)
    feat_scores = tf.zeros(shape, dtype=dtype)

    feat_ymin = tf.zeros(shape, dtype=dtype)
    feat_xmin = tf.zeros(shape, dtype=dtype)
    feat_ymax = tf.ones(shape, dtype=dtype)
    feat_xmax = tf.ones(shape, dtype=dtype)

    def jaccard_with_anchors(bbox):
        """Compute jaccard score between a box and the anchors.
        """
        int_ymin = tf.maximum(ymin, bbox[0])
        int_xmin = tf.maximum(xmin, bbox[1])
        int_ymax = tf.minimum(ymax, bbox[2])
        int_xmax = tf.minimum(xmax, bbox[3])
        h = tf.maximum(int_ymax - int_ymin, 0.)
        w = tf.maximum(int_xmax - int_xmin, 0.)
        # Volumes
        inter_vol = h * w
        union_vol = vol_anchors - inter_vol \
            + (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
        jaccard = tf.div(inter_vol, union_vol)
        return jaccard

    def intersection_with_anchors(bbox):
        """Compute intersection between score a box and the anchors.
        """
        int_ymin = tf.maximum(ymin, bbox[0])
        int_xmin = tf.maximum(xmin, bbox[1])
        int_ymax = tf.minimum(ymax, bbox[2])
        int_xmax = tf.minimum(xmax, bbox[3])
        h = tf.maximum(int_ymax - int_ymin, 0.)
        w = tf.maximum(int_xmax - int_xmin, 0.)
        inter_vol = h * w
        scores = tf.div(inter_vol, vol_anchors)
        return scores

    def condition(i, feat_labels, feat_scores,
                  feat_ymin, feat_xmin, feat_ymax, feat_xmax):
        """Condition: check label index.
        """
        r = tf.less(i, tf.shape(labels))
        return r[0]

    def body(i, feat_labels, feat_scores,
             feat_ymin, feat_xmin, feat_ymax, feat_xmax):
        """Body: update feature labels, scores and bboxes.
        Follow the original SSD paper for that purpose:
          - assign values when jaccard > 0.5;
          - only update if beat the score of other bboxes.
        """
        # Jaccard score
        label = labels[i]
        bbox = bboxes[i]
        jaccard = jaccard_with_anchors(bbox)

        # Mask: check threshold + scores + no annotations + num_classes.
        mask = tf.greater(jaccard, feat_scores)
        mask = tf.logical_and(mask, feat_scores > -0.5)

        # [TODO] : Fix the following line
        mask = tf.logical_and(mask, label < num_classes)
        imask = tf.cast(mask, tf.int64)
        fmask = tf.cast(mask, dtype)
        
        # Update values using mask.
        feat_labels = imask * label + (1 - imask) * feat_labels
        feat_scores = tf.where(mask, jaccard, feat_scores)

        feat_ymin = fmask * bbox[0] + (1 - fmask) * feat_ymin
        feat_xmin = fmask * bbox[1] + (1 - fmask) * feat_xmin
        feat_ymax = fmask * bbox[2] + (1 - fmask) * feat_ymax
        feat_xmax = fmask * bbox[3] + (1 - fmask) * feat_xmax

        return [i+1, feat_labels, feat_scores,
                feat_ymin, feat_xmin, feat_ymax, feat_xmax]

    # Main loop definition.
    i = 0
    [i, feat_labels, feat_scores,
     feat_ymin, feat_xmin,
     feat_ymax, feat_xmax] = tf.while_loop(condition, body,
                                           [i, feat_labels, feat_scores,
                                            feat_ymin, feat_xmin,
                                            feat_ymax, feat_xmax])
    # Transform to center / size.
    feat_cy = (feat_ymax + feat_ymin) / 2.
    feat_cx = (feat_xmax + feat_xmin) / 2.
    feat_h = feat_ymax - feat_ymin
    feat_w = feat_xmax - feat_xmin
    # Encode features.
    feat_cy = (feat_cy - yref) / href / prior_scaling[0]
    feat_cx = (feat_cx - xref) / wref / prior_scaling[1]
    feat_h = tf.log(feat_h / href) / prior_scaling[2]
    feat_w = tf.log(feat_w / wref) / prior_scaling[3]
    # Use SSD ordering: x / y / w / h instead of ours.
    feat_localizations = tf.stack([feat_cx, feat_cy, feat_w, feat_h], axis=-1)
    return feat_labels, feat_localizations, feat_scores

def get_shape(x, rank=None):
    """Returns the dimensions of a Tensor as list of integers or scale tensors.
    Args:
      x: N-d Tensor;
      rank: Rank of the Tensor. If None, will try to guess it.
    Returns:
      A list of `[d1, d2, ..., dN]` corresponding to the dimensions of the
        input tensor.  Dimensions that are statically known are python integers,
        otherwise they are integer scalar tensors.
    """
    if x.get_shape().is_fully_defined():
        return x.get_shape().as_list()
    else:
        static_shape = x.get_shape()
        if rank is None:
            static_shape = static_shape.as_list()
            rank = len(static_shape)
        else:
            static_shape = x.get_shape().with_rank(rank).as_list()
        dynamic_shape = tf.unstack(tf.shape(x), rank)
        return [s if s is not None else d
                for s, d in zip(static_shape, dynamic_shape)]

def tf_ssd_bboxes_select_layer(predictions_layer, localizations_layer,
                               num_classes,
                               select_threshold=0.5,
                               ignore_class=0,
                               scope=None):
    """Extract classes, scores and bounding boxes from features in one layer.
    Batch-compatible: inputs are supposed to have batch-type shapes.
    Args:
      predictions_layer: A SSD prediction layer;
      localizations_layer: A SSD localization layer;
      select_threshold: Classification threshold for selecting a box. All boxes
        under the threshold are set to 'zero'. If None, no threshold applied.
    Return:
      d_scores, d_bboxes: Dictionary of scores and bboxes Tensors of
        size Batches X N x 1 | 4. Each key corresponding to a class.
    """
    select_threshold = 0.0 if select_threshold is None else select_threshold
    with tf.name_scope(scope, 'ssd_bboxes_select_layer', [predictions_layer, localizations_layer]):
        # Reshape features: Batches x N x N_labels | 4
        p_shape = get_shape(predictions_layer)
        predictions_layer = tf.reshape(predictions_layer, tf.stack([p_shape[0], -1, p_shape[-1]]))
        l_shape = get_shape(localizations_layer)
        localizations_layer = tf.reshape(localizations_layer, tf.stack([l_shape[0], -1, l_shape[-1]]))
        d_scores = {}
        d_bboxes = {}
        for c in range(0, num_classes):
            if c != ignore_class:
                # Remove boxes under the threshold.
                scores = predictions_layer[:, :, c]
                fmask = tf.cast(tf.greater_equal(scores, select_threshold), scores.dtype)
                scores = scores * fmask
                bboxes = localizations_layer * tf.expand_dims(fmask, axis=-1)
                # Append to dictionary.
                d_scores[c] = scores
                d_bboxes[c] = bboxes

        return d_scores, d_bboxes


def tf_ssd_bboxes_select(predictions_net, localizations_net,
                         num_classes,
                         select_threshold=0.5,
                         ignore_class=0,
                         scope=None):
    """Extract classes, scores and bounding boxes from network output layers.
    Batch-compatible: inputs are supposed to have batch-type shapes.
    Args:
      predictions_net: List of SSD prediction layers;
      localizations_net: List of localization layers;
      select_threshold: Classification threshold for selecting a box. All boxes
        under the threshold are set to 'zero'. If None, no threshold applied.
    Return:
      d_scores, d_bboxes: Dictionary of scores and bboxes Tensors of
        size Batches X N x 1 | 4. Each key corresponding to a class.
    """
    with tf.name_scope(scope, 'ssd_bboxes_select',
                       [predictions_net, localizations_net]):
        l_scores = []
        l_bboxes = []
        for i in range(len(predictions_net)):
            predictions = tf.nn.softmax(predictions_net[i][0])
            scores, bboxes = tf_ssd_bboxes_select_layer(predictions,
                                                        localizations_net[i],
                                                        num_classes,
                                                        select_threshold,
                                                        ignore_class)
            l_scores.append(scores)
            l_bboxes.append(bboxes)
        # Concat results.
        d_scores = {}
        d_bboxes = {}
        for c in l_scores[0].keys():
            ls = [s[c] for s in l_scores]
            lb = [b[c] for b in l_bboxes]
            d_scores[c] = tf.concat(ls, axis=1)
            d_bboxes[c] = tf.concat(lb, axis=1)
        return d_scores, d_bboxes

def bboxes_sort(scores, bboxes, top_k=400, scope=None):
    """Sort bounding boxes by decreasing order and keep only the top_k.
    If inputs are dictionnaries, assume every key is a different class.
    Assume a batch-type input.
    Args:
      scores: Batch x N Tensor/Dictionary containing float scores.
      bboxes: Batch x N x 4 Tensor/Dictionary containing boxes coordinates.
      top_k: Top_k boxes to keep.
    Return:
      scores, bboxes: Sorted Tensors/Dictionaries of shape Batch x Top_k x 1|4.
    """
    # Dictionaries as inputs.
    if isinstance(scores, dict) or isinstance(bboxes, dict):
        with tf.name_scope(scope, 'bboxes_sort_dict'):
            d_scores = {}
            d_bboxes = {}
            for c in scores.keys():
                s, b = bboxes_sort(scores[c], bboxes[c], top_k=top_k)
                d_scores[c] = s
                d_bboxes[c] = b
            return d_scores, d_bboxes

    # Tensors inputs.
    with tf.name_scope(scope, 'bboxes_sort', [scores, bboxes]):
        # Sort scores...
        scores, idxes = tf.nn.top_k(scores, k=top_k, sorted=True)

        # Trick to be able to use tf.gather: map for each element in the first dim.
        def fn_gather(bboxes, idxes):
            bb = tf.gather(bboxes, idxes)
            return [bb]
        r = tf.map_fn(lambda x: fn_gather(x[0], x[1]), 
                      [bboxes, idxes],
                      dtype=[bboxes.dtype],
                      parallel_iterations=10,
                      back_prop=False,
                      swap_memory=False,
                      infer_shape=True)
        bboxes = r[0]
        return scores, bboxes

def pad_axis(x, offset, size, axis=0, name=None):
    """Pad a tensor on an axis, with a given offset and output size.
    The tensor is padded with zero (i.e. CONSTANT mode). Note that the if the
    `size` is smaller than existing size + `offset`, the output tensor
    was the latter dimension.
    Args:
      x: Tensor to pad;
      offset: Offset to add on the dimension chosen;
      size: Final size of the dimension.
    Return:
      Padded tensor whose dimension on `axis` is `size`, or greater if
      the input vector was larger.
    """
    with tf.name_scope(name, 'pad_axis'):
        shape = get_shape(x)
        rank = len(shape)
        # Padding description.
        new_size = tf.maximum(size-offset-shape[axis], 0)
        pad1 = tf.stack([0]*axis + [offset] + [0]*(rank-axis-1))
        pad2 = tf.stack([0]*axis + [new_size] + [0]*(rank-axis-1))
        paddings = tf.stack([pad1, pad2], axis=1)
        x = tf.pad(x, paddings, mode='CONSTANT')
        # Reshape, to get fully defined shape if possible.
        # TODO: fix with tf.slice
        shape[axis] = size
        x = tf.reshape(x, tf.stack(shape))
        return x

def bboxes_nms(scores, bboxes, nms_threshold=0.5, keep_top_k=200, scope=None):
    """Apply non-maximum selection to bounding boxes. In comparison to TF
    implementation, use classes information for matching.
    Should only be used on single-entries. Use batch version otherwise.
    Args:
      scores: N Tensor containing float scores.
      bboxes: N x 4 Tensor containing boxes coordinates.
      nms_threshold: Matching threshold in NMS algorithm;
      keep_top_k: Number of total object to keep after NMS.
    Return:
      classes, scores, bboxes Tensors, sorted by score.
        Padded with zero if necessary.
    """
    with tf.name_scope(scope, 'bboxes_nms_single', [scores, bboxes]):
        # Apply NMS algorithm.
        idxes = tf.image.non_max_suppression(bboxes, scores,
                                             keep_top_k, nms_threshold)
        scores = tf.gather(scores, idxes)
        bboxes = tf.gather(bboxes, idxes)
        # Pad results.
        scores = pad_axis(scores, 0, keep_top_k, axis=0)
        bboxes = pad_axis(bboxes, 0, keep_top_k, axis=0)
        return scores, bboxes


def bboxes_nms_batch(scores, bboxes, nms_threshold=0.5, keep_top_k=200,
                     scope=None):
    """Apply non-maximum selection to bounding boxes. In comparison to TF
    implementation, use classes information for matching.
    Use only on batched-inputs. Use zero-padding in order to batch output
    results.
    Args:
      scores: Batch x N Tensor/Dictionary containing float scores.
      bboxes: Batch x N x 4 Tensor/Dictionary containing boxes coordinates.
      nms_threshold: Matching threshold in NMS algorithm;
      keep_top_k: Number of total object to keep after NMS.
    Return:
      scores, bboxes Tensors/Dictionaries, sorted by score.
        Padded with zero if necessary.
    """
    # Dictionaries as inputs.
    if isinstance(scores, dict) or isinstance(bboxes, dict):
        with tf.name_scope(scope, 'bboxes_nms_batch_dict'):
            d_scores = {}
            d_bboxes = {}
            for c in scores.keys():
                s, b = bboxes_nms_batch(scores[c], bboxes[c],
                                        nms_threshold=nms_threshold,
                                        keep_top_k=keep_top_k)
                d_scores[c] = s
                d_bboxes[c] = b
            return d_scores, d_bboxes

    # Tensors inputs.
    with tf.name_scope(scope, 'bboxes_nms_batch'):
        r = tf.map_fn(lambda x: bboxes_nms(x[0], x[1],
                                           nms_threshold, keep_top_k),
                      (scores, bboxes),
                      dtype=(scores.dtype, bboxes.dtype),
                      parallel_iterations=10,
                      back_prop=False,
                      swap_memory=False,
                      infer_shape=True)
        scores, bboxes = r
        return scores, bboxes

def bboxes_clip(bbox_ref, bboxes, scope=None):
    """Clip bounding boxes to a reference box.
    Batch-compatible if the first dimension of `bbox_ref` and `bboxes`
    can be broadcasted.
    Args:
      bbox_ref: Reference bounding box. Nx4 or 4 shaped-Tensor;
      bboxes: Bounding boxes to clip. Nx4 or 4 shaped-Tensor or dictionary.
    Return:
      Clipped bboxes.
    """
    # Bboxes is dictionary.
    if isinstance(bboxes, dict):
        with tf.name_scope(scope, 'bboxes_clip_dict'):
            d_bboxes = {}
            for c in bboxes.keys():
                d_bboxes[c] = bboxes_clip(bbox_ref, bboxes[c])
            return d_bboxes

    # Tensors inputs.
    with tf.name_scope(scope, 'bboxes_clip'):
        # Easier with transposed bboxes. Especially for broadcasting.
        bbox_ref = tf.transpose(bbox_ref)
        bboxes = tf.transpose(bboxes)
        # Intersection bboxes and reference bbox.
        ymin = tf.maximum(bboxes[0], bbox_ref[0])
        xmin = tf.maximum(bboxes[1], bbox_ref[1])
        ymax = tf.minimum(bboxes[2], bbox_ref[2])
        xmax = tf.minimum(bboxes[3], bbox_ref[3])
        # Double check! Empty boxes when no-intersection.
        ymin = tf.minimum(ymin, ymax)
        xmin = tf.minimum(xmin, xmax)
        bboxes = tf.transpose(tf.stack([ymin, xmin, ymax, xmax], axis=0))
        return bboxes

def decode_predictions(overall_predictions, overall_anchors, num_classes, clipping_bbox=None, select_threshold=None, nms_threshold=0.5, top_k=400, keep_top_k=200, prior_scaling=[0.1, 0.1, 0.2, 0.2]):
    """ Decode the boxes given by the network back to the image domain """
    bboxes = []
    for index, (predictions, anchors) in enumerate(zip(overall_predictions, overall_anchors)):
        yref, xref, href, wref = anchors
        pred_cx = predictions[1][:, :, :, :, 0] * wref * prior_scaling[1] + xref
        pred_cy = predictions[1][:, :, :, :, 1] * href * prior_scaling[0] + yref
        # pdb.set_trace()
        pred_w = wref * tf.exp(predictions[1][:, :, :, :, 2] * prior_scaling[3])
        pred_h = href * tf.exp(predictions[1][:, :, :, :, 3] * prior_scaling[2])
	
        xmin = pred_cx - pred_w / 2.
        ymin = pred_cy - pred_h / 2.
        xmax = pred_cx + pred_w / 2.
        ymax = pred_cy + pred_h / 2.
        bboxes.append(tf.stack([ymin, xmin, ymax, xmax], axis=-1))

    rscores, rbboxes = tf_ssd_bboxes_select(overall_predictions, bboxes, num_classes, select_threshold)
    rscores, rbboxes = bboxes_sort(rscores, rbboxes, top_k=top_k)
    rscores, rbboxes = bboxes_nms_batch(rscores, rbboxes, nms_threshold=nms_threshold, keep_top_k=keep_top_k)
    rbboxes = bboxes_clip(clipping_bbox, rbboxes)

    return rscores, rbboxes

def channel_to_last(inputs,
                    data_format='NHWC',
                    scope=None):
    """Move the channel axis to the last dimension. Allows to
    provide a single output format whatever the input data format.
    Args:
      inputs: Input Tensor;
      data_format: NHWC or NCHW.
    Return:
      Input in NHWC format.
    """
    with tf.name_scope(scope, 'channel_to_last', [inputs]):
        if data_format == 'NHWC':
            net = inputs
        elif data_format == 'NCHW':
            net = tf.transpose(inputs, perm=(0, 2, 3, 1))
        return net

def tensor_shape(x, rank=3):
    """Returns the dimensions of a tensor.
    Args:
      image: A N-D Tensor of shape.
    Returns:
      A list of dimensions. Dimensions that are statically known are python
        integers,otherwise they are integer scalar tensors.
    """
    if x.get_shape().is_fully_defined():
        return x.get_shape().as_list()
    else:
        static_shape = x.get_shape().with_rank(rank).as_list()
        dynamic_shape = tf.unstack(tf.shape(x), rank)
        return [s if s is not None else d
                for s, d in zip(static_shape, dynamic_shape)]

def overlay_bboxes(bboxes, tf_image):
    """ This function draws the bounding boxes on a batch of images """
    tf_image  = tf.image.draw_bounding_boxes(tf_image, tf.expand_dims(bboxes, axis=1), name="overlay_bboxes")

    return tf_image

def adapt_overlay_bboxes_pred(bboxes, tf_image):
    """ This function draws the bounding boxes on a batch of images """
    new_bboxes_list = [bboxes[:, 0], bboxes[:, 1], tf.exp(bboxes[:, 2]) + bboxes[:, 0], tf.exp(bboxes[:, 3]) + bboxes[:, 1]]
    bboxes = tf.stack(new_bboxes_list, axis=1)
    tf_image  = tf.image.draw_bounding_boxes(tf_image, tf.expand_dims(bboxes, axis=1), name="overlay_bboxes")
    return tf_image

def overlay_bboxes_eval(detection_scores, detection_bboxes, tf_image):
    """ This function draws the bounding boxes on a batch of images """
    for class_index in detection_bboxes:
        tf_image  = tf.image.draw_bounding_boxes(tf_image, detection_bboxes[class_index], name="overlay_bboxes")

    return tf_image

def overlay_bboxes_ground_truth(gt_classes, gt_bboxes, tf_image, batch_size=1):
    """ This function draws the bounding boxes on a batch of images """
    draw_bbox_tensor = []
    for batch_ind in range(batch_size):
        draw_bbox_tensor.append(tf.image.draw_bounding_boxes(tf.expand_dims(tf_image[batch_ind], axis=0), tf.expand_dims(gt_bboxes[batch_ind], axis=0), name="overlay_bboxes"))        
    tf_image = tf.concat(draw_bbox_tensor, axis=0)
    return tf_image


 
def _phase_shift(X, r):
    # Helper function with main phase shift operation
    return tf.depth_to_space(X, r)

def pixel_shuffler(X, scale, channels, activation=tf.identity, name=None):
    # Main OP that you can arbitrarily use in you tensorflow code
    if channels>1:
        Xc = tf.split(X, channels//(scale**2), axis=3)
        X = tf.concat([_phase_shift(x, scale) for x in Xc], axis=3)
    return activation(X, name=name)
#def _phase_shift(I, r):
#    # Helper function with main phase shift operation
#    bsize, a, b, c = I.get_shape().as_list()
#    X = tf.reshape(I, [-1, a, b, r, r])
#    X = tf.transpose(X, (0, 1, 2, 4, 3))  # bsize, a, b, 1, 1
#    X = tf.split(X, a, axis=1)  # a, [bsize, b, r, r]
#    X = tf.concat([tf.squeeze(x) for x in X], 2)  # bsize, b, a*r, r
#    X = tf.split(X, b, axis=1)  # b, [bsize, a*r, r]
#    X = tf.concat([tf.squeeze(x) for x in X], 2)  #bsize, a*r, b*r
#    return tf.reshape(X, [-1, a*r, b*r, 1])
#
#def pixel_shuffler(X, scale, channels, activation=tf.identity, name=None):
#    # Main OP that you can arbitrarily use in you tensorflow code
#    if channels>1:
#        Xc = tf.split(X, channels//(scale**2), axis=3)
#        X = tf.concat([_phase_shift(x, scale) for x in Xc], axis=3)
#    else:
#        X = _phase_shift(X, r)
#    return activation(X, name=name)