Merge pull request #6 from srihari-humbarwadi/paste_masks

Paste masks using bilinear grid sampling

official/vision/beta/projects/panoptic_maskrcnn/modeling/layers/panoptic_segmentation_generator.py

@@ -18,6 +18,7 @@ from typing import List
 
 import tensorflow as tf
 
+from official.vision.beta.projects.panoptic_maskrcnn.modeling.layers import paste_masks
 
 class PanopticSegmentationGenerator(tf.keras.layers.Layer):
   """Panoptic segmentation generator layer."""
@@ -79,34 +80,10 @@ class PanopticSegmentationGenerator(tf.keras.layers.Layer):
     }
     super(PanopticSegmentationGenerator, self).__init__(**kwargs)
 
-  def _paste_mask(self, box, mask):
-    pasted_mask = tf.ones(
-        self._output_size + [1], dtype=mask.dtype) * self._void_class_label
-
-    ymin = tf.clip_by_value(box[0], 0, self._output_size[0])
-    xmin = tf.clip_by_value(box[1], 0, self._output_size[1])
-    ymax = tf.clip_by_value(box[2] + 1, 0, self._output_size[0])
-    xmax = tf.clip_by_value(box[3] + 1, 0, self._output_size[1])
-    box_height = ymax - ymin
-    box_width = xmax - xmin
-
-    if not (box_height == 0 or box_width == 0):
-      # resize mask to match the shape of the instance bounding box
-      resized_mask = tf.image.resize(
-          mask,
-          size=(box_height, box_width),
-          method='bilinear')
-      resized_mask = tf.cast(resized_mask, dtype=mask.dtype)
-
-      # paste resized mask on a blank mask that matches image shape
-      pasted_mask = tf.raw_ops.TensorStridedSliceUpdate(
-          input=pasted_mask,
-          begin=[ymin, xmin],
-          end=[ymax, xmax],
-          strides=[1, 1],
-          value=resized_mask)
-
-    return pasted_mask
+  def build(self, input_shape):
+    grid_sampler = paste_masks.BilinearGridSampler(align_corners=False)
+    self._paste_masks_fn = paste_masks.PasteMasks(
+        output_size=self._output_size, grid_sampler=grid_sampler)
 
   def _generate_panoptic_masks(self, boxes, scores, classes, detections_masks,
                                segmentation_mask):
@@ -138,6 +115,8 @@ class PanopticSegmentationGenerator(tf.keras.layers.Layer):
         - category_mask: A `tf.Tensor` for category masks.
         - instance_mask: A `tf.Tensor for instance masks.
     """
+    # Paste instance masks
+
     # Offset stuff class predictions
     segmentation_mask = tf.where(
         tf.logical_or(
@@ -167,6 +146,10 @@ class PanopticSegmentationGenerator(tf.keras.layers.Layer):
       loop_end_idx = tf.minimum(
           tf.cast(loop_end_idx, dtype=tf.int32),
           self._max_num_detections)
+    
+      pasted_masks = self._paste_masks_fn((
+          detections_masks[:loop_end_idx],
+          boxes[:loop_end_idx]))
 
       # add things segmentation to panoptic masks
       for i in range(loop_end_idx):
@@ -174,9 +157,7 @@ class PanopticSegmentationGenerator(tf.keras.layers.Layer):
         # the overlaps are resolved based on confidence score
         instance_idx = sorted_indices[i]
 
-        pasted_mask = self._paste_mask(
-            box=boxes[instance_idx],
-            mask=detections_masks[instance_idx])
+        pasted_mask = pasted_masks[instance_idx]
 
         class_id = tf.cast(classes[instance_idx], dtype=tf.float32)
 
@@ -248,7 +229,7 @@ class PanopticSegmentationGenerator(tf.keras.layers.Layer):
 
     batched_scores = detections['detection_scores']
     batched_classes = detections['detection_classes']
-    batched_boxes = tf.cast(detections['detection_boxes'], dtype=tf.int32)
+    batched_boxes = detections['detection_boxes']
     batched_detections_masks = tf.expand_dims(
         detections['detection_masks'], axis=-1)
 

official/vision/beta/projects/panoptic_maskrcnn/modeling/layers/paste_masks.py

+import tensorflow as tf
+
+
+class BilinearGridSampler(tf.keras.layers.Layer):
+  def __init__(self, align_corners, **kwargs):
+    super(BilinearGridSampler, self).__init__(**kwargs)
+    self.align_corners = align_corners
+
+    self._config = {
+        'align_corners': align_corners
+    }
+
+  def build(self, input_shape):
+    features_shape, _, _ = input_shape
+    _, height, width, channels = features_shape.as_list()
+
+    self._height = height
+    self._width = width
+    self._channels = channels
+
+  def _valid_coordinates(self, x_coord, y_coord):
+    return tf.logical_and(
+        tf.logical_and(
+            tf.greater_equal(x_coord, 0),
+            tf.greater_equal(y_coord, 0)),
+        tf.logical_and(
+            tf.less(x_coord, self._width),
+            tf.less(y_coord, self._width)))
+
+  def _get_pixel(self, features, x_coord, y_coord):
+    x_coord = tf.cast(x_coord, dtype=tf.int32)
+    y_coord = tf.cast(y_coord, dtype=tf.int32)
+
+    clipped_x = tf.clip_by_value(x_coord, 0, self._width - 1)
+    clipped_y = tf.clip_by_value(y_coord, 0, self._height - 1)
+
+    batch_size, _, _, _ = features.shape.as_list()
+    if batch_size is None:
+      batch_size = tf.shape(features)[0]
+
+    batch_indices = tf.reshape(
+        tf.range(batch_size, dtype=tf.int32),
+        shape=[batch_size, 1, 1])
+    batch_indices = tf.tile(
+        batch_indices,
+        multiples=[1, x_coord.shape[1], x_coord.shape[2]])
+    indices = tf.cast(
+        tf.stack([batch_indices, clipped_y, clipped_x], axis=-1),
+        dtype=tf.int32)
+    gathered_pixels = tf.gather_nd(features, indices)
+
+    return tf.where(
+        tf.expand_dims(self._valid_coordinates(x_coord, y_coord), axis=-1),
+        gathered_pixels,
+        tf.zeros_like(gathered_pixels))
+
+  def call(self, inputs):
+    features, x_coord, y_coord = inputs
+
+    x_coord += 1
+    y_coord += 1
+
+    if self.align_corners:
+      x_coord = (x_coord * 0.5) * (self._width - 1)
+      y_coord = (y_coord * 0.5) * (self._height - 1)
+    else:
+      x_coord = (x_coord * self._width - 1) * 0.5
+      y_coord = (y_coord * self._height - 1) * 0.5
+
+    left = tf.floor(x_coord)
+    top = tf.floor(y_coord)
+    right = left + 1
+    bottom = top + 1
+
+    top_left = (right - x_coord) * (bottom - y_coord)
+    top_right = (x_coord - left) * (bottom - y_coord)
+    bottom_left = (right - x_coord) * (y_coord - top)
+    bottom_right = (x_coord - left) * (y_coord - top)
+
+    i_top_left = self._get_pixel(features, left, top)
+    i_top_right = self._get_pixel(features, right, top)
+    i_bottom_left = self._get_pixel(features, left, bottom)
+    i_bottom_right = self._get_pixel(features, right, bottom)
+
+    i_top_left *= tf.expand_dims(top_left, axis=-1)
+    i_top_right *= tf.expand_dims(top_right, axis=-1)
+    i_bottom_left *= tf.expand_dims(bottom_left, axis=-1)
+    i_bottom_right *= tf.expand_dims(bottom_right, axis=-1)
+
+    interpolated_features = tf.math.add_n(
+        [i_top_left, i_top_right, i_bottom_left, i_bottom_right])
+    return interpolated_features
+
+  def get_config(self):
+    return self._config_dict
+
+
+
+class PasteMasks(tf.keras.layers.Layer):
+
+  def __init__(self, output_size, grid_sampler, **kwargs):
+    super(PasteMasks, self).__init__(**kwargs)
+    self._output_size = output_size
+    self._grid_sampler = grid_sampler
+
+    self._config = {
+        'output_size': output_size,
+        'grid_sampler': grid_sampler
+    }
+
+  def build(self, input_shape):
+    self._x_coords = tf.range(0, self._output_size[1], dtype=tf.float32)
+    self._y_coords = tf.range(0, self._output_size[0], dtype=tf.float32)
+
+  def call(self, inputs):
+    masks, boxes = inputs
+    y0, x0, y1, x1 = tf.split(boxes, 4, axis=1)
+
+    x_coords = tf.cast(self._x_coords, dtype=boxes.dtype)
+    y_coords = tf.cast(self._y_coords, dtype=boxes.dtype)
+    x_coords = (x_coords - x0) / (x1 - x0) * 2 - 1
+    y_coords = (y_coords - y0) / (y1 - y0) * 2 - 1
+
+    x_coords = tf.tile(
+        tf.expand_dims(x_coords, axis=1),
+        multiples=[1, self._output_size[0], 1])
+    y_coords = tf.tile(
+        tf.expand_dims(y_coords, axis=2),
+        multiples=[1, 1, self._output_size[1]])
+    pasted_masks = self._grid_sampler((masks, x_coords, y_coords))
+    return pasted_masks
+
+  def get_config(self):
+    return self._config_dict