Fix bugs for multilevel_native_crop_and_resize

research/object_detection/utils/spatial_transform_ops.py

@@ -414,29 +414,23 @@ def multilevel_roi_align(features, boxes, box_levels, output_size,
 
 def multilevel_native_crop_and_resize(images, boxes, box_levels, 
                                       crop_size, scope=None):
-  #FIXME: fix docstring
-  """doc string."""
+  """Multilevel native crop and resize.
+
+  Same as `multilevel_matmul_crop_and_resize` but uses tf.image.crop_and_resize.
+  """
   if box_levels is None:
     return native_crop_and_resize(images[0], boxes, crop_size, scope=None)
-  croped_feature_list = []
+  cropped_feature_list = []
   for level, image in enumerate(images):
-    # indicies = tf.boolean_mask(tf.range(0, boxes[0]), box_levels == level)
-    # print(indicies)
-    # level_boxes = tf.gather(boxes, indicies)
-    # print(level_boxes)
-    # level_boxes = boxes[box_levels == level]
-    # level_boxes = tf.reshape(level_boxes,
-    #                          [1, -1] + level_boxes.shape.as_list()[1:])
+    # For each level, crop the feature according to all boxes
+    # set the cropped feature not at this level to 0 tensor.
     cropped = native_crop_and_resize(image, boxes, crop_size)
-    print(cropped)
-    cond = tf.tile(tf.equal(box_levels, level)[:, :, tf.newaxis], [1, 1] + [tf.math.reduce_prod(cropped.shape.as_list()[2:])])
+    cond = tf.tile(tf.equal(box_levels, level)[:, :, tf.newaxis],
+              [1, 1] + [tf.math.reduce_prod(cropped.shape.as_list()[2:])])
     cond = tf.reshape(cond, cropped.shape)
-    print(cond)
     cropped_final = tf.where(cond, cropped, tf.zeros_like(cropped))
-    # cropped[tf.where(box_levels != level)] = tf.zeros(crop_size)
-    print(cropped_final)
-    croped_feature_list.append(cropped_final)
-  return tf.concat(croped_feature_list, axis=0)
+    cropped_feature_list.append(cropped_final)
+  return tf.math.reduce_sum(cropped_feature_list, axis=0)
 
 
 def native_crop_and_resize(image, boxes, crop_size, scope=None):
@@ -461,8 +455,36 @@ def native_crop_and_resize(image, boxes, crop_size, scope=None):
 
 def multilevel_matmul_crop_and_resize(images, boxes, box_levels, crop_size, 
                                       extrapolation_value=0.0, scope=None):
-  #FIXME: fix docstring
-  """doc string."""
+  """Multilevel native crop and resize.
+
+  Same as `matmul_crop_and_resize` but crop images according to box levels.
+
+  Args:
+    images: A list of 4-D tensor of shape 
+      [batch, image_height, image_width, depth] representing features of 
+      different size.
+    boxes: A `Tensor` of type `float32` or 'bfloat16'.
+      A 3-D tensor of shape `[batch, num_boxes, 4]`. The boxes are specified in
+      normalized coordinates and are of the form `[y1, x1, y2, x2]`. A
+      normalized coordinate value of `y` is mapped to the image coordinate at
+      `y * (image_height - 1)`, so as the `[0, 1]` interval of normalized image
+      height is mapped to `[0, image_height - 1] in image height coordinates.
+      We do allow y1 > y2, in which case the sampled crop is an up-down flipped
+      version of the original image. The width dimension is treated similarly.
+      Normalized coordinates outside the `[0, 1]` range are allowed, in which
+      case we use `extrapolation_value` to extrapolate the input image values.
+    box_levels: A 2-D tensor of shape [batch, num_boxes] representing the level
+      of the box.
+    crop_size: A list of two integers `[crop_height, crop_width]`. All
+      cropped image patches are resized to this size. The aspect ratio of the
+      image content is not preserved. Both `crop_height` and `crop_width` need
+      to be positive.
+    extrapolation_value: a float value to use for extrapolation.
+    scope: A name for the operation (optional).
+  
+  Returns:
+    A 5-D tensor of shape `[batch, num_boxes, crop_height, crop_width, depth]`
+  """
   with tf.name_scope(scope, 'MatMulCropAndResize'):
     if box_levels is None:
       box_levels = tf.zeros(tf.shape(boxes)[:2], dtype=tf.int32)