Merge pull request #5 from srihari-humbarwadi/training

 * set segmentation loss weight to 0.5
 * skip boxes with zero area while pasting masks

Author: srihari-humbarwadi

official/vision/beta/projects/panoptic_maskrcnn/configs/panoptic_maskrcnn.py

@@ -107,7 +107,7 @@ class Losses(maskrcnn.Losses):
   semantic_segmentation_use_groundtruth_dimension: bool = True
   semantic_segmentation_top_k_percent_pixels: float = 1.0
   instance_segmentation_weight: float = 1.0
-  semantic_segmentation_weight: float = 1.0
+  semantic_segmentation_weight: float = 0.5
 
 
 @dataclasses.dataclass
@@ -170,7 +170,8 @@ def panoptic_fpn_coco() -> cfg.ExperimentConfig:
     is_thing.append(True if idx <= num_thing_categories else False)
 
   config = cfg.ExperimentConfig(
-      runtime=cfg.RuntimeConfig(mixed_precision_dtype='bfloat16'),
+    runtime=cfg.RuntimeConfig(
+      mixed_precision_dtype='bfloat16', enable_xla=True),
       task=PanopticMaskRCNNTask(
           init_checkpoint='gs://cloud-tpu-checkpoints/vision-2.0/resnet50_imagenet/ckpt-28080',  # pylint: disable=line-too-long
           init_checkpoint_modules=['backbone'],

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

@@ -79,26 +79,27 @@ def _paste_mask(self, box, mask):
     pasted_mask = tf.ones(
         self._output_size + [1], dtype=mask.dtype) * self._void_class_label
 
-    ymin = box[0]
-    xmin = box[1]
+    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
 
-    # resize mask to match the shape of the instance bounding box
-    resized_mask = tf.image.resize(
-        mask,
-        size=(box_height, box_width),
-        method='nearest')
-
-    # 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)
+    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='nearest')
+
+      # 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