Implement thresholding instance masks with a configurable value.

PiperOrigin-RevId: 351890707

official/vision/beta/configs/maskrcnn.py

@@ -31,11 +31,13 @@ from official.vision.beta.configs import decoders
 @dataclasses.dataclass
 class TfExampleDecoder(hyperparams.Config):
   regenerate_source_id: bool = False
+  mask_binarize_threshold: Optional[float] = None
 
 
 @dataclasses.dataclass
 class TfExampleDecoderLabelMap(hyperparams.Config):
   regenerate_source_id: bool = False
+  mask_binarize_threshold: Optional[float] = None
   label_map: str = ''
 
 

official/vision/beta/dataloaders/tf_example_decoder.py

@@ -34,7 +34,8 @@ class TfExampleDecoder(decoder.Decoder):
 
   def __init__(self,
                include_mask=False,
-               regenerate_source_id=False):
+               regenerate_source_id=False,
+               mask_binarize_threshold=None):
     self._include_mask = include_mask
     self._regenerate_source_id = regenerate_source_id
     self._keys_to_features = {
@@ -50,6 +51,7 @@ class TfExampleDecoder(decoder.Decoder):
         'image/object/area': tf.io.VarLenFeature(tf.float32),
         'image/object/is_crowd': tf.io.VarLenFeature(tf.int64),
     }
+    self._mask_binarize_threshold = mask_binarize_threshold
     if include_mask:
       self._keys_to_features.update({
           'image/object/mask': tf.io.VarLenFeature(tf.string),
@@ -151,6 +153,9 @@ class TfExampleDecoder(decoder.Decoder):
     if self._include_mask:
       masks = self._decode_masks(parsed_tensors)
 
+      if self._mask_binarize_threshold is not None:
+        masks = tf.cast(masks > self._mask_binarize_threshold, tf.float32)
+
     decoded_tensors = {
         'source_id': source_id,
         'image': image,

official/vision/beta/dataloaders/tf_example_label_map_decoder.py

@@ -27,9 +27,11 @@ from official.vision.beta.dataloaders import tf_example_decoder
 class TfExampleDecoderLabelMap(tf_example_decoder.TfExampleDecoder):
   """Tensorflow Example proto decoder."""
 
-  def __init__(self, label_map, include_mask=False, regenerate_source_id=False):
+  def __init__(self, label_map, include_mask=False, regenerate_source_id=False,
+               mask_binarize_threshold=None):
     super(TfExampleDecoderLabelMap, self).__init__(
-        include_mask=include_mask, regenerate_source_id=regenerate_source_id)
+        include_mask=include_mask, regenerate_source_id=regenerate_source_id,
+        mask_binarize_threshold=mask_binarize_threshold)
     self._keys_to_features.update({
         'image/object/class/text': tf.io.VarLenFeature(tf.string),
     })

official/vision/beta/tasks/maskrcnn.py

@@ -110,12 +110,14 @@ class MaskRCNNTask(base_task.Task):
     if params.decoder.type == 'simple_decoder':
       decoder = tf_example_decoder.TfExampleDecoder(
           include_mask=self._task_config.model.include_mask,
-          regenerate_source_id=decoder_cfg.regenerate_source_id)
+          regenerate_source_id=decoder_cfg.regenerate_source_id,
+          mask_binarize_threshold=decoder_cfg.mask_binarize_threshold)
     elif params.decoder.type == 'label_map_decoder':
       decoder = tf_example_label_map_decoder.TfExampleDecoderLabelMap(
           label_map=decoder_cfg.label_map,
           include_mask=self._task_config.model.include_mask,
-          regenerate_source_id=decoder_cfg.regenerate_source_id)
+          regenerate_source_id=decoder_cfg.regenerate_source_id,
+          mask_binarize_threshold=decoder_cfg.mask_binarize_threshold)
     else:
       raise ValueError('Unknown decoder type: {}!'.format(params.decoder.type))