add evaluate_tflite for object detection in Model Maker.

PiperOrigin-RevId: 362245976

Author: ziyeqinghan

tensorflow_examples/lite/model_maker/core/task/model_spec/object_detector_spec.py

@@ -16,6 +16,7 @@
 import collections
 import os
 import tempfile
+from typing import Tuple, Dict
 
 from absl import logging
 import tensorflow as tf
@@ -26,6 +27,7 @@
 from tensorflow_examples.lite.model_maker.third_party.efficientdet import hparams_config
 from tensorflow_examples.lite.model_maker.third_party.efficientdet import utils
 from tensorflow_examples.lite.model_maker.third_party.efficientdet.keras import efficientdet_keras
+from tensorflow_examples.lite.model_maker.third_party.efficientdet.keras import eval_tflite
 from tensorflow_examples.lite.model_maker.third_party.efficientdet.keras import inference
 from tensorflow_examples.lite.model_maker.third_party.efficientdet.keras import label_util
 from tensorflow_examples.lite.model_maker.third_party.efficientdet.keras import postprocess
@@ -202,8 +204,10 @@ def train(self,
         validation_steps=validation_steps)
     return model
 
-  def evaluate(self, model, dataset, steps, json_file=None):
-    """Evaluate the EfficientDet keras model."""
+  def _get_evaluator_and_label_map(
+      self, json_file: str
+  ) -> Tuple[coco_metric.EvaluationMetric, collections.OrderedDict]:
+    """Gets evaluator and label_map for evaluation."""
     label_map = label_util.get_label_map(self.config.label_map)
     # Sorts label_map.keys since pycocotools.cocoeval uses sorted catIds
     # (category ids) in COCOeval class.
@@ -213,6 +217,43 @@ def evaluate(self, model, dataset, steps, json_file=None):
         filename=json_file, label_map=label_map)
 
     evaluator.reset_states()
+    return evaluator, label_map
+
+  def _get_metric_dict(self, evaluator: coco_metric.EvaluationMetric,
+                       label_map: collections.OrderedDict) -> Dict[str, float]:
+    """Gets the metric dict for evaluation."""
+    metrics = evaluator.result()
+    metric_dict = {}
+    for i, name in enumerate(evaluator.metric_names):
+      metric_dict[name] = metrics[i]
+
+    if label_map:
+      for i, cid in enumerate(label_map.keys()):
+        name = 'AP_/%s' % label_map[cid]
+        metric_dict[name] = metrics[i + len(evaluator.metric_names)]
+    return metric_dict
+
+  def evaluate(self,
+               model: tf.keras.Model,
+               dataset: tf.data.Dataset,
+               steps: int,
+               json_file: str = None) -> Dict[str, float]:
+    """Evaluate the EfficientDet keras model.
+
+    Args:
+      model: The keras model to be evaluated.
+      dataset: tf.data.Dataset used for evaluation.
+      steps: Number of steps to evaluate the model.
+      json_file: JSON with COCO data format containing golden bounding boxes.
+        Used for validation. If None, use the ground truth from the dataloader.
+        Refer to
+        https://towardsdatascience.com/coco-data-format-for-object-detection-a4c5eaf518c5
+          for the description of COCO data format.
+
+    Returns:
+      A dict contains AP metrics.
+    """
+    evaluator, label_map = self._get_evaluator_and_label_map(json_file)
     dataset = dataset.take(steps)
 
     @tf.function
@@ -228,18 +269,66 @@ def _get_detections(images, labels):
       ], [])
 
     dataset = self.ds_strategy.experimental_distribute_dataset(dataset)
-    for (images, labels) in dataset:
+    progbar = tf.keras.utils.Progbar(steps)
+    for i, (images, labels) in enumerate(dataset):
       self.ds_strategy.run(_get_detections, (images, labels))
+      progbar.update(i)
 
-    metrics = evaluator.result()
-    metric_dict = {}
-    for i, name in enumerate(evaluator.metric_names):
-      metric_dict[name] = metrics[i]
+    metric_dict = self._get_metric_dict(evaluator, label_map)
+    return metric_dict
 
-    if label_map:
-      for i, cid in enumerate(label_map.keys()):
-        name = 'AP_/%s' % label_map[cid]
-        metric_dict[name] = metrics[i + len(evaluator.metric_names)]
+  def evaluate_tflite(self,
+                      tflite_filepath: str,
+                      dataset: tf.data.Dataset,
+                      steps: int,
+                      json_file: str = None) -> Dict[str, float]:
+    """Evaluate the EfficientDet TFLite model.
+
+    Args:
+      tflite_filepath: File path to the TFLite model.
+      dataset: tf.data.Dataset used for evaluation.
+      steps: Number of steps to evaluate the model.
+      json_file: JSON with COCO data format containing golden bounding boxes.
+        Used for validation. If None, use the ground truth from the dataloader.
+        Refer to
+        https://towardsdatascience.com/coco-data-format-for-object-detection-a4c5eaf518c5
+          for the description of COCO data format.
+
+    Returns:
+      A dict contains AP metrics.
+    """
+    # TODO(b/182441458): Use the task library for evaluation instead once it
+    # supports python interface.
+    evaluator, label_map = self._get_evaluator_and_label_map(json_file)
+    dataset = dataset.take(steps)
+
+    lite_runner = eval_tflite.LiteRunner(tflite_filepath, only_network=False)
+    progbar = tf.keras.utils.Progbar(steps)
+    for i, (images, labels) in enumerate(dataset):
+      # Get the output result after post-processing NMS op.
+      nms_boxes, nms_classes, nms_scores, _ = lite_runner.run(images)
+
+      # CLASS_OFFSET is used since label_id for `background` is 0 in label_map
+      # while it's not actually included the model. We don't need to add the
+      # offset in the Android application.
+      nms_classes += postprocess.CLASS_OFFSET
+
+      height, width = utils.parse_image_size(self.config.image_size)
+      normalize_factor = tf.constant([height, width, height, width],
+                                     dtype=tf.float32)
+      nms_boxes *= normalize_factor
+      if labels['image_scales'] is not None:
+        scales = tf.expand_dims(tf.expand_dims(labels['image_scales'], -1), -1)
+        nms_boxes = nms_boxes * tf.cast(scales, nms_boxes.dtype)
+      detections = postprocess.generate_detections_from_nms_output(
+          nms_boxes, nms_classes, nms_scores, labels['source_ids'])
+
+      detections = postprocess.transform_detections(detections)
+      evaluator.update_state(labels['groundtruth_data'].numpy(),
+                             detections.numpy())
+      progbar.update(i)
+
+    metric_dict = self._get_metric_dict(evaluator, label_map)
     return metric_dict
 
   def export_saved_model(self,

tensorflow_examples/lite/model_maker/core/task/object_detector.py

@@ -15,9 +15,11 @@
 
 import os
 import tempfile
+from typing import Dict
 
 import tensorflow as tf
 from tensorflow_examples.lite.model_maker.core import compat
+from tensorflow_examples.lite.model_maker.core.data_util import object_detector_dataloader
 from tensorflow_examples.lite.model_maker.core.export_format import ExportFormat
 from tensorflow_examples.lite.model_maker.core.task import custom_model
 from tensorflow_examples.lite.model_maker.core.task import model_spec as ms
@@ -161,6 +163,14 @@ def evaluate(self, data, batch_size=None):
     return self.model_spec.evaluate(self.model, ds, steps,
                                     data.annotations_json_file)
 
+  def evaluate_tflite(
+      self, tflite_filepath: str,
+      data: object_detector_dataloader.DataLoader) -> Dict[str, float]:
+    """Evaluate the TFLite model."""
+    ds = data.gen_dataset(self.model_spec, batch_size=1, is_training=False)
+    return self.model_spec.evaluate_tflite(tflite_filepath, ds, len(data),
+                                           data.annotations_json_file)
+
   def _export_saved_model(self, saved_model_dir):
     """Saves the model to Tensorflow SavedModel."""
     self.model_spec.export_saved_model(saved_model_dir)

tensorflow_examples/lite/model_maker/core/task/object_detector_test.py

@@ -75,6 +75,11 @@ def testEfficientDetLite0(self):
         'efficientdet_lite0_metadata.json')
     self.assertTrue(filecmp.cmp(json_output_file, expected_json_file))
 
+    # Evaluate the TFLite model.
+    task.evaluate_tflite(output_path, data)
+    self.assertIsInstance(metrics, dict)
+    self.assertGreaterEqual(metrics['AP'], 0)
+
     # Export the model to quantized TFLite model.
     # TODO(b/175173304): Skips the test for stable tensorflow 2.4 for now since
     # it fails. Will revert this change after TF upgrade.