1. add missing part in ann_json_dict in object_detector_dataloader.

2. Move create_pascal_voc function to test_util so that other unittest (object_detector task unittest) can call it as well.

PiperOrigin-RevId: 351260755

Author: ziyeqinghan

tensorflow_examples/lite/model_maker/core/data_util/object_detector_dataloader.py

@@ -53,7 +53,7 @@ def _get_object_detector_cache_filenames(cache_dir, image_dir, annotations_dir,
   """Gets cache filenames for obejct detector."""
   if cache_dir is None:
     cache_dir = tempfile.mkdtemp()
-    print('Create the cache directory: %s.', cache_dir)
+    print('Create the cache directory: %s.' % cache_dir)
   cache_prefix = _get_cache_prefix(image_dir, annotations_dir, annotations_list)
   cache_prefix = os.path.join(cache_dir, cache_prefix)
 
@@ -211,7 +211,12 @@ def _write_pascal_tfrecord(cls, images_dir, annotations_dir, label_map_dict,
     for idx, name in label_map_dict.items():
       label_name2id_dict[name] = idx
     writers = [tf.io.TFRecordWriter(path) for path in tfrecord_files]
+
     ann_json_dict = {'images': [], 'annotations': [], 'categories': []}
+    for class_id, class_name in label_map_dict.items():
+      c = {'supercategory': 'none', 'id': class_id, 'name': class_name}
+      ann_json_dict['categories'].append(c)
+
     # Gets the paths to annotations.
     if annotations_list:
       ann_path_list = [
@@ -246,7 +251,7 @@ def _write_pascal_tfrecord(cls, images_dir, annotations_dir, label_map_dict,
       writer.close()
 
     with tf.io.gfile.GFile(annotations_json_file, 'w') as f:
-      json.dump(ann_json_dict, f)
+      json.dump(ann_json_dict, f, indent=2)
 
   def gen_dataset(self,
                   model_spec,

tensorflow_examples/lite/model_maker/core/data_util/object_detector_dataloader_test.py

@@ -16,10 +16,11 @@
 from __future__ import division
 from __future__ import print_function
 
+import filecmp
 import os
 
 import numpy as np
-import PIL.Image
+
 import tensorflow as tf
 from tensorflow_examples.lite.model_maker.core import test_util
 
@@ -40,28 +41,9 @@ def __init__(self, model_name):
 
 class ObjectDectectorDataLoaderTest(tf.test.TestCase):
 
-  def _create_pascal_voc(self):
-    # Saves the image into images_dir.
-    image_file_name = '2012_12.jpg'
-    image_data = np.random.rand(256, 256, 3)
-    images_dir = os.path.join(self.get_temp_dir(), 'images')
-    os.mkdir(images_dir)
-    save_path = os.path.join(images_dir, image_file_name)
-    image = PIL.Image.fromarray(image_data, 'RGB')
-    image.save(save_path)
-
-    # Gets the annonation path.
-    annotations_path = test_util.get_test_data_path('2012_12.xml')
-    annotations_dir = os.path.dirname(annotations_path)
-
-    label_map = {
-        1: 'person',
-        2: 'notperson',
-    }
-    return images_dir, annotations_dir, label_map
-
   def test_from_pascal_voc(self):
-    images_dir, annotations_dir, label_map = self._create_pascal_voc()
+    images_dir, annotations_dir, label_map = test_util.create_pascal_voc(
+        self.get_temp_dir())
     model_spec = MockDetectorModelSpec('efficientdet-lite0')
 
     data = object_detector_dataloader.DataLoader.from_pascal_voc(
@@ -71,6 +53,11 @@ def test_from_pascal_voc(self):
     self.assertLen(data, 1)
     self.assertEqual(data.label_map, label_map)
 
+    self.assertTrue(os.path.isfile(data.annotations_json_file))
+    self.assertGreater(os.path.getsize(data.annotations_json_file), 0)
+    expected_json_file = test_util.get_test_data_path('annotations.json')
+    self.assertTrue(filecmp.cmp(data.annotations_json_file, expected_json_file))
+
     ds = data.gen_dataset(model_spec, batch_size=1, is_training=False)
     for i, (images, labels) in enumerate(ds):
       self.assertEqual(i, 0)

tensorflow_examples/lite/model_maker/core/data_util/testdata/annotations.json

@@ -0,0 +1,39 @@
+{
+  "images": [
+    {
+      "file_name": "2012_12.jpg",
+      "height": 256,
+      "width": 256,
+      "id": 1
+    }
+  ],
+  "annotations": [
+    {
+      "area": 16384,
+      "iscrowd": 0,
+      "image_id": 1,
+      "bbox": [
+        64,
+        64,
+        128,
+        128
+      ],
+      "category_id": 1,
+      "id": 1,
+      "ignore": 0,
+      "segmentation": []
+    }
+  ],
+  "categories": [
+    {
+      "supercategory": "none",
+      "id": 1,
+      "name": "person"
+    },
+    {
+      "supercategory": "none",
+      "id": 2,
+      "name": "notperson"
+    }
+  ]
+}

tensorflow_examples/lite/model_maker/core/test_util.py

@@ -19,9 +19,11 @@
 import functools
 import os
 import shutil
+import tempfile
 
 from absl import flags
 import numpy as np
+import PIL.Image
 
 import tensorflow.compat.v2 as tf
 from tensorflow_examples.lite.model_maker.core import compat
@@ -138,6 +140,31 @@ def get_dataloader(data_size, input_shape, num_classes, max_input_value=1000):
   return data
 
 
+def create_pascal_voc(temp_dir=None):
+  """Creates test data with PASCAL VOC format."""
+  if temp_dir is None or not tf.io.gfile.exists(temp_dir):
+    temp_dir = tempfile.mkdtemp()
+
+  # Saves the image into images_dir.
+  image_file_name = "2012_12.jpg"
+  image_data = np.random.rand(256, 256, 3)
+  images_dir = os.path.join(temp_dir, "images")
+  os.mkdir(images_dir)
+  save_path = os.path.join(images_dir, image_file_name)
+  image = PIL.Image.fromarray(image_data, "RGB")
+  image.save(save_path)
+
+  # Gets the annonation path.
+  annotations_path = get_test_data_path("2012_12.xml")
+  annotations_dir = os.path.dirname(annotations_path)
+
+  label_map = {
+      1: "person",
+      2: "notperson",
+  }
+  return images_dir, annotations_dir, label_map
+
+
 def is_same_output(tflite_file,
                    keras_model,
                    input_tensors,