Merge pull request #8 from srihari-humbarwadi/rescale_predictions

rescale predictions to original size during evaluation

official/vision/beta/evaluation/panoptic_quality_evaluator.py

@@ -32,11 +32,34 @@ import tensorflow as tf
 from official.vision.beta.evaluation import panoptic_quality
 
 
+def rescale_masks(groundtruth, prediction, image_info):
+  rescale_size = tf.cast(
+      tf.math.ceil(image_info[1, :] / image_info[2, :]), tf.int32)
+  image_shape = tf.cast(image_info[0, :], tf.int32)
+  offsets = tf.cast(image_info[3, :], tf.int32)
+
+  prediction = tf.image.resize(
+      tf.expand_dims(prediction, axis=-1),
+      rescale_size,
+      method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
+  prediction = tf.image.crop_to_bounding_box(
+      prediction,
+      offsets[0], offsets[1],
+      image_shape[0],
+      image_shape[1])
+  groundtruth = tf.image.crop_to_bounding_box(
+      tf.expand_dims(groundtruth, axis=-1), 0, 0,
+      image_shape[0], image_shape[1])
+
+  return (
+    tf.expand_dims(groundtruth[:, :, 0], axis=0),
+    tf.expand_dims(prediction[:, :, 0], axis=0))
+
 class PanopticQualityEvaluator:
   """Panoptic Quality metric class."""
 
   def __init__(self, num_categories, ignored_label, max_instances_per_category,
-               offset, is_thing=None):
+               offset, is_thing=None, rescale_predictions=False):
     """Constructs Panoptic Quality evaluation class.
 
     The class provides the interface to Panoptic Quality metrics_fn.
@@ -55,10 +78,14 @@ class PanopticQualityEvaluator:
         `is_thing[category_id]` is True iff that category is a "thing" category
         instead of "stuff." Default to `None`, and it means categories are not
         classified into these two categories.
+      rescale_predictions: `bool`, whether to scale back prediction to original
+        image sizes. If True, y_true['image_info'] is used to rescale
+        predictions.        
     """
     self._pq_metric_module = panoptic_quality.PanopticQuality(
         num_categories, ignored_label, max_instances_per_category, offset)
     self._is_thing = is_thing
+    self._rescale_predictions = rescale_predictions
     self._required_prediction_fields = ['category_mask', 'instance_mask']
     self._required_groundtruth_fields = ['category_mask', 'instance_mask']
     self.reset_states()
@@ -110,6 +137,13 @@ class PanopticQualityEvaluator:
         Required fields:
           - category_mask: a numpy array of uint16 of shape [batch_size, H, W].
           - instance_mask: a numpy array of uint16 of shape [batch_size, H, W].
+          - image_info: [batch, 4, 2], a tensor that holds information about
+          original and preprocessed images. Each entry is in the format of
+          [[original_height, original_width], [input_height, input_width],
+          [y_scale, x_scale], [y_offset, x_offset]], where [desired_height,
+          desired_width] is the actual scaled image size, and [y_scale, x_scale]
+          is the scaling factor, which is the ratio of scaled dimension /
+          original dimension.
       predictions: a dictionary of tensors including the fields below. See
         different parsers under `../dataloader` for more details.
         Required fields:
@@ -132,4 +166,29 @@ class PanopticQualityEvaluator:
         raise ValueError(
             'Missing the required key `{}` in groundtruths!'.format(k))
 
-    self._pq_metric_module.compare_and_accumulate(groundtruths, predictions)
+    if self._rescale_predictions:
+      for idx in range(len(groundtruths['category_mask'])):
+        image_info = groundtruths['image_info'][idx]
+        groundtruth_category_mask, prediction_category_mask = rescale_masks(
+            groundtruths['category_mask'][idx],
+            predictions['category_mask'][idx],
+            image_info)
+        groundtruth_instance_mask, prediction_instance_mask = rescale_masks(
+            groundtruths['instance_mask'][idx],
+            predictions['instance_mask'][idx],
+            image_info)
+        _groundtruths = {
+            'category_mask': groundtruth_category_mask,
+            'instance_mask': groundtruth_instance_mask
+            }
+        _predictions = {
+            'category_mask': prediction_category_mask,
+            'instance_mask': prediction_instance_mask
+            }
+        _groundtruths, _predictions = self._convert_to_numpy(
+            _groundtruths, _predictions)
+
+        self._pq_metric_module.compare_and_accumulate(
+            _groundtruths, _predictions)
+    else:
+      self._pq_metric_module.compare_and_accumulate(groundtruths, predictions)

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

@@ -121,6 +121,7 @@ class PanopticQualityEvaluator(hyperparams.Config):
   offset: int = 256 * 256 * 256
   is_thing: List[float] = dataclasses.field(
       default_factory=list)
+  rescale_predictions: bool = False
   report_per_class_metrics: bool = False
 
 
@@ -203,6 +204,9 @@ def panoptic_fpn_coco() -> cfg.ExperimentConfig:
               input_path=os.path.join(_COCO_INPUT_PATH_BASE, 'val*'),
               is_training=False,
               global_batch_size=eval_batch_size,
+              parser=Parser(
+                  segmentation_resize_eval_groundtruth=False,
+                  segmentation_groundtruth_padded_size=[640, 640]),
               drop_remainder=False),
           annotation_file=os.path.join(_COCO_INPUT_PATH_BASE,
                                        'instances_val2017.json'),
@@ -211,7 +215,8 @@ def panoptic_fpn_coco() -> cfg.ExperimentConfig:
           panoptic_quality_evaluator=PanopticQualityEvaluator(
               num_categories=num_panoptic_categories,
               ignored_label=0,
-              is_thing=is_thing)),
+              is_thing=is_thing,
+              rescale_predictions=True)),
       trainer=cfg.TrainerConfig(
           train_steps=22500,
           validation_steps=validation_steps,

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

@@ -259,7 +259,8 @@ class PanopticMaskRCNNTask(maskrcnn.MaskRCNNTask):
             ignored_label=pq_config.ignored_label,
             max_instances_per_category=pq_config.max_instances_per_category,
             offset=pq_config.offset,
-            is_thing=pq_config.is_thing)
+            is_thing=pq_config.is_thing,
+            rescale_predictions=pq_config.rescale_predictions)
 
     return metrics
 
@@ -383,7 +384,8 @@ class PanopticMaskRCNNTask(maskrcnn.MaskRCNNTask):
           'category_mask':
               labels['groundtruths']['gt_panoptic_category_mask'],
           'instance_mask':
-              labels['groundtruths']['gt_panoptic_instance_mask']
+              labels['groundtruths']['gt_panoptic_instance_mask'],
+          'image_info': labels['image_info']
             }
       logs.update({
           self.panoptic_quality_metric.name: