提取带有填充的矩形框

Question

I'm trying to extract values From a 2d Tensor inside multiple rectangular regions. I want to crop rectangular regions while setting all values outside the box to zero.

For example from the 9 x 9 image I want to get two separate images with values inside the two rectangular red boxes, while setting the rest of the values to zero. Is there a convenient way to do this with tensorflow slicing?

One way I thought of approaching this is defining a mask array that is 1 inside the box and 0 outside and multiply it with the input array. But this requires looping over the number of boxes, each time changing which values of the mask are set to 0. Is there a faster and more efficient way to do this without using for loops? Is there an equivalent of crop and replace function in tensorflow? Here's the code I have with the for loop. Appreciate any input on this. Thanks

import tensorflow as tf
import matplotlib.pyplot as plt
import matplotlib.patches as patches

tf.reset_default_graph()

size = 9 # size of input image
num_boxes = 2 # number of rectangular boxes


def get_cutout(X, bboxs):
    """Returns copies of X with values only inside bboxs"""
    out = []
    for i in range(num_boxes):
        bbox = bboxs[i] # get rectangular box coordinates
        Y = tf.Variable(np.zeros((size, size)), dtype=tf.float32) # define temporary mask
        # set values of mask inside box to 1
        t = [Y[bbox[0]:bbox[2], bbox[2]:bbox[3]].assign(
            tf.ones((bbox[2]-bbox[0], bbox[3]-bbox[2])))]
        with tf.control_dependencies(t):
            mask = tf.identity(Y) 
        out.append(X * mask) # get values inside rectangular box
    return out, X

#define a 9x9 input image X and convert to tensor
in_x = np.eye(size)
in_x[0:3]=np.random.rand(3,9)
X = tf.constant(in_x , dtype=tf.float32)

bboxs = tf.placeholder(tf.int32, [None, 4]) # placeholder for rectangular box

X_outs = get_cutout(X, bboxs)

# coordintes of box ((bottom left x, bottom left y, top right x, top right y))
in_bbox = [[1,3,3,6], [4,3,7,8]] 
feed_dict = {bboxs: in_bbox}

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    x_out= sess.run(X_outs, feed_dict=feed_dict)

# plot results
vmin = np.min(x_out[2])
vmax = np.max(x_out[2])
fig, ax = plt.subplots(nrows=1, ncols=1+len(in_bbox),figsize=(10,2))
im = ax[0].imshow(x_out[2], vmin=vmin, vmax=vmax, origin='lower')
plt.colorbar(im, ax=ax[0])
ax[0].set_title("input X")
for i, bbox in enumerate(in_bbox):
    bottom_left = (bbox[2]-0.5, bbox[0]-0.5)
    width = bbox[3]-bbox[2]
    height = bbox[2]- bbox[0]
    rect = patches.Rectangle(bottom_left, width, height,
                             linewidth=1,edgecolor='r',facecolor='none')
    ax[0].add_patch(rect)
    ax[i+1].set_title("extract values in box {}".format(i+1))
    im = ax[i + 1].imshow(x_out[0][i], vmin=vmin, vmax=vmax, origin='lower')
    plt.colorbar(im,ax=ax[i+1])

Answer 1

The mask can be created using tf.pad .

 crop = tf.ones([3, 3])
 # "before_axis_x" how many padding will be added before cropping zone over the axis x
 # "after_axis_x" how many padding will be added after cropping zone over the axis x
 mask = tf.pad(crop, [[before_axis_0, after_axis_0], [before_axis_1, after_axis_1]]

 tf.mask(image, mask) # creates the extracted image

To have the same behavior as tf.image.crop_and_resize, here is a function that will take an array of boxes and will return an array of extracted images with padding.

def extract_with_padding(image, boxes):
  """
   boxes: tensor of shape [num_boxes, 4]. 
          boxes are the coordinates of the extracted part
          box is an array [y1, x1, y2, x2] 
          where [y1, x1] (respectively [y2, x2]) are the coordinates 
          of the top left (respectively bottom right ) part of the image
   image: tensor containing the initial image
  """
  extracted = []
  shape = tf.shape(image)
  for b in boxes:
    crop = tf.ones([3, 3])

    mask = tf.pad(crop, [[b[0], shape[0] - b[2]], [b[1] , shape[1] - b[3]]])
    extracted.append(tf.boolean_mask(image, mask))

  return extracted

Answer 2

Thanks for that really nice function @edkevekeh. I've had to modify it slightly to get it to do what I want. One, I couldn't iterate over boxes which is a Tensor object. Plus, the crop size is determined by the box and not always 3x3. Also, tf.boolean_mask returns the crop, but I want to keep the crop, but replace outside the crop with 0. So I replaced the tf.boolean_mask with multiplication.

For my use case num_boxes can be large, so I wanted to know if there was a more efficient way than a for loop, guess not. My modified version of @edkevekeh's solution if anyone else needs it.

def extract_with_padding(image, boxes):
    """
    boxes: tensor of shape [num_boxes, 4]. 
          boxes are the coordinates of the extracted part
          box is an array [y1, x1, y2, x2] 
          where [y1, x1] (respectively [y2, x2]) are the coordinates 
          of the top left (respectively bottom right ) part of the image
    image: tensor containing the initial image
    """
    extracted = []
    shape = tf.shape(image)
    for i in range(boxes.shape[0]):
        b = boxes[i]
        crop = tf.ones([b[2] - b[0], b[3] - b[1]])
        mask = tf.pad(crop, [[b[0], shape[0] - b[2]], [b[1] , shape[1] - b[3]]])
        extracted.append(image*mask)
    return extracted