detection_test.py 5.73 KB
Newer Older
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
1
# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
Abdullah Rashwan's avatar
Abdullah Rashwan committed
2
3
4
5
6
7
8
9
10
11
12
13
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
Yeqing Li's avatar
Yeqing Li committed
14

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
15
"""Test for image detection export lib."""
Abdullah Rashwan's avatar
Abdullah Rashwan committed
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

import io
import os

from absl.testing import parameterized
import numpy as np
from PIL import Image
import tensorflow as tf

from official.common import registry_imports  # pylint: disable=unused-import
from official.core import exp_factory
from official.vision.beta.serving import detection


class DetectionExportTest(tf.test.TestCase, parameterized.TestCase):

Fan Yang's avatar
Fan Yang committed
32
  def _get_detection_module(self, experiment_name, input_type):
Abdullah Rashwan's avatar
Abdullah Rashwan committed
33
34
    params = exp_factory.get_exp_config(experiment_name)
    params.task.model.backbone.resnet.model_id = 18
Xianzhi Du's avatar
Xianzhi Du committed
35
    params.task.model.detection_generator.nms_version = 'batched'
Abdullah Rashwan's avatar
Abdullah Rashwan committed
36
    detection_module = detection.DetectionModule(
Fan Yang's avatar
Fan Yang committed
37
38
39
40
        params,
        batch_size=1,
        input_image_size=[640, 640],
        input_type=input_type)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
41
42
    return detection_module

Hongkun Yu's avatar
Hongkun Yu committed
43
44
45
46
  def _export_from_module(self, module, input_type, save_directory):
    signatures = module.get_inference_signatures(
        {input_type: 'serving_default'})
    tf.saved_model.save(module, save_directory, signatures=signatures)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
47

A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
48
  def _get_dummy_input(self, input_type, batch_size, image_size):
Abdullah Rashwan's avatar
Abdullah Rashwan committed
49
    """Get dummy input for the given input type."""
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
50
    h, w = image_size
Abdullah Rashwan's avatar
Abdullah Rashwan committed
51
52

    if input_type == 'image_tensor':
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
53
      return tf.zeros((batch_size, h, w, 3), dtype=np.uint8)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
54
    elif input_type == 'image_bytes':
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
55
      image = Image.fromarray(np.zeros((h, w, 3), dtype=np.uint8))
Abdullah Rashwan's avatar
Abdullah Rashwan committed
56
57
58
59
      byte_io = io.BytesIO()
      image.save(byte_io, 'PNG')
      return [byte_io.getvalue() for b in range(batch_size)]
    elif input_type == 'tf_example':
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
60
      image_tensor = tf.zeros((h, w, 3), dtype=tf.uint8)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
61
62
63
64
65
66
67
68
69
      encoded_jpeg = tf.image.encode_jpeg(tf.constant(image_tensor)).numpy()
      example = tf.train.Example(
          features=tf.train.Features(
              feature={
                  'image/encoded':
                      tf.train.Feature(
                          bytes_list=tf.train.BytesList(value=[encoded_jpeg])),
              })).SerializeToString()
      return [example for b in range(batch_size)]
Fan Yang's avatar
Fan Yang committed
70
71
    elif input_type == 'tflite':
      return tf.zeros((batch_size, h, w, 3), dtype=np.float32)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
72
73

  @parameterized.parameters(
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
74
75
76
      ('image_tensor', 'fasterrcnn_resnetfpn_coco', [384, 384]),
      ('image_bytes', 'fasterrcnn_resnetfpn_coco', [640, 640]),
      ('tf_example', 'fasterrcnn_resnetfpn_coco', [640, 640]),
Fan Yang's avatar
Fan Yang committed
77
      ('tflite', 'fasterrcnn_resnetfpn_coco', [640, 640]),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
78
79
80
      ('image_tensor', 'maskrcnn_resnetfpn_coco', [640, 640]),
      ('image_bytes', 'maskrcnn_resnetfpn_coco', [640, 384]),
      ('tf_example', 'maskrcnn_resnetfpn_coco', [640, 640]),
Fan Yang's avatar
Fan Yang committed
81
      ('tflite', 'maskrcnn_resnetfpn_coco', [640, 640]),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
82
83
84
      ('image_tensor', 'retinanet_resnetfpn_coco', [640, 640]),
      ('image_bytes', 'retinanet_resnetfpn_coco', [640, 640]),
      ('tf_example', 'retinanet_resnetfpn_coco', [384, 640]),
Fan Yang's avatar
Fan Yang committed
85
      ('tflite', 'retinanet_resnetfpn_coco', [640, 640]),
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
86
87
88
      ('image_tensor', 'retinanet_resnetfpn_coco', [384, 384]),
      ('image_bytes', 'retinanet_spinenet_coco', [640, 640]),
      ('tf_example', 'retinanet_spinenet_coco', [640, 384]),
Fan Yang's avatar
Fan Yang committed
89
      ('tflite', 'retinanet_spinenet_coco', [640, 640]),
Abdullah Rashwan's avatar
Abdullah Rashwan committed
90
  )
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
91
  def test_export(self, input_type, experiment_name, image_size):
Abdullah Rashwan's avatar
Abdullah Rashwan committed
92
    tmp_dir = self.get_temp_dir()
Fan Yang's avatar
Fan Yang committed
93
    module = self._get_detection_module(experiment_name, input_type)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
94

Hongkun Yu's avatar
Hongkun Yu committed
95
    self._export_from_module(module, input_type, tmp_dir)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
96
97

    self.assertTrue(os.path.exists(os.path.join(tmp_dir, 'saved_model.pb')))
Hongkun Yu's avatar
Hongkun Yu committed
98
99
100
101
102
103
    self.assertTrue(
        os.path.exists(os.path.join(tmp_dir, 'variables', 'variables.index')))
    self.assertTrue(
        os.path.exists(
            os.path.join(tmp_dir, 'variables',
                         'variables.data-00000-of-00001')))
Abdullah Rashwan's avatar
Abdullah Rashwan committed
104
105

    imported = tf.saved_model.load(tmp_dir)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
106
    detection_fn = imported.signatures['serving_default']
Abdullah Rashwan's avatar
Abdullah Rashwan committed
107

Hongkun Yu's avatar
Hongkun Yu committed
108
109
    images = self._get_dummy_input(
        input_type, batch_size=1, image_size=image_size)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
110

Fan Yang's avatar
Fan Yang committed
111
112
113
114
115
116
117
118
    if input_type == 'tflite':
      processed_images = tf.zeros(image_size + [3], dtype=tf.float32)
      anchor_boxes = module._build_anchor_boxes()
      image_info = tf.convert_to_tensor(
          [image_size, image_size, [1.0, 1.0], [0, 0]], dtype=tf.float32)
    else:
      processed_images, anchor_boxes, image_info = module._build_inputs(
          tf.zeros((224, 224, 3), dtype=tf.uint8))
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
119
    image_shape = image_info[1, :]
Abdullah Rashwan's avatar
Abdullah Rashwan committed
120
    image_shape = tf.expand_dims(image_shape, 0)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
121
    processed_images = tf.expand_dims(processed_images, 0)
Abdullah Rashwan's avatar
Abdullah Rashwan committed
122
123
124
    for l, l_boxes in anchor_boxes.items():
      anchor_boxes[l] = tf.expand_dims(l_boxes, 0)

Hongkun Yu's avatar
Hongkun Yu committed
125
    expected_outputs = module.model(
Abdullah Rashwan's avatar
Abdullah Rashwan committed
126
127
128
129
        images=processed_images,
        image_shape=image_shape,
        anchor_boxes=anchor_boxes,
        training=False)
A. Unique TensorFlower's avatar
A. Unique TensorFlower committed
130
    outputs = detection_fn(tf.constant(images))
Abdullah Rashwan's avatar
Abdullah Rashwan committed
131
132
133
134

    self.assertAllClose(outputs['num_detections'].numpy(),
                        expected_outputs['num_detections'].numpy())

Fan Yang's avatar
Fan Yang committed
135
136
137
138
139
140
141
  def test_build_model_fail_with_none_batch_size(self):
    params = exp_factory.get_exp_config('retinanet_resnetfpn_coco')
    with self.assertRaisesRegex(
        ValueError, 'batch_size cannot be None for detection models.'):
      detection.DetectionModule(
          params, batch_size=None, input_image_size=[640, 640])

Hongkun Yu's avatar
Hongkun Yu committed
142

Abdullah Rashwan's avatar
Abdullah Rashwan committed
143
144
if __name__ == '__main__':
  tf.test.main()