Internal change

PiperOrigin-RevId: 276210315

official/vision/image_classification/imagenet_preprocessing.py

@@ -54,7 +54,7 @@ _SHUFFLE_BUFFER = 10000
 _R_MEAN = 123.68
 _G_MEAN = 116.78
 _B_MEAN = 103.94
-_CHANNEL_MEANS = [_R_MEAN, _G_MEAN, _B_MEAN]
+CHANNEL_MEANS = [_R_MEAN, _G_MEAN, _B_MEAN]
 
 # The lower bound for the smallest side of the image for aspect-preserving
 # resizing. For example, if an image is 500 x 1000, it will be resized to
@@ -524,4 +524,4 @@ def preprocess_image(image_buffer, bbox, output_height, output_width,
 
   image.set_shape([output_height, output_width, num_channels])
 
-  return _mean_image_subtraction(image, _CHANNEL_MEANS, num_channels)
+  return _mean_image_subtraction(image, CHANNEL_MEANS, num_channels)

official/vision/image_classification/resnet_model.py

@@ -32,7 +32,7 @@ from tensorflow.python.keras import initializers
 from tensorflow.python.keras import layers
 from tensorflow.python.keras import models
 from tensorflow.python.keras import regularizers
-
+from official.vision.image_classification import imagenet_preprocessing
 
 L2_WEIGHT_DECAY = 1e-4
 BATCH_NORM_DECAY = 0.9
@@ -223,28 +223,40 @@ def conv_block(input_tensor,
 
 def resnet50(num_classes,
              batch_size=None,
-             use_l2_regularizer=True):
+             use_l2_regularizer=True,
+             rescale_inputs=False):
   """Instantiates the ResNet50 architecture.
 
   Args:
     num_classes: `int` number of classes for image classification.
     batch_size: Size of the batches for each step.
     use_l2_regularizer: whether to use L2 regularizer on Conv/Dense layer.
+    rescale_inputs: whether to rescale inputs from 0 to 1.
 
   Returns:
       A Keras model instance.
   """
   input_shape = (224, 224, 3)
   img_input = layers.Input(shape=input_shape, batch_size=batch_size)
+  if rescale_inputs:
+    # Hub image modules expect inputs in the range [0, 1]. This rescales these
+    # inputs to the range expected by the trained model.
+    x = layers.Lambda(
+        lambda x: x * 255.0 - backend.constant(
+            imagenet_preprocessing.CHANNEL_MEANS,
+            shape=[1, 1, 3],
+            dtype=x.dtype),
+        name='rescale')(
+            img_input)
+  else:
+    x = img_input
 
   if backend.image_data_format() == 'channels_first':
     x = layers.Lambda(
         lambda x: backend.permute_dimensions(x, (0, 3, 1, 2)),
-        name='transpose')(
-            img_input)
+        name='transpose')(x)
     bn_axis = 1
   else:  # channels_last
-    x = img_input
     bn_axis = 3
 
   x = layers.ZeroPadding2D(padding=(3, 3), name='conv1_pad')(x)

official/vision/image_classification/tfhub_export.py

@@ -26,8 +26,8 @@ from absl import flags
 
 import tensorflow as tf
 
-from official.vision.image_classification import resnet_model
 from official.vision.image_classification import imagenet_preprocessing
+from official.vision.image_classification import resnet_model
 
 FLAGS = flags.FLAGS
 
@@ -39,9 +39,11 @@ flags.DEFINE_string("export_path", None,
 
 def export_tfhub(model_path, hub_destination):
   """Restores a tf.keras.Model and saves for TF-Hub."""
-  model = resnet_model.resnet50(num_classes=imagenet_preprocessing.NUM_CLASSES)
+  model = resnet_model.resnet50(
+      num_classes=imagenet_preprocessing.NUM_CLASSES, rescale_inputs=True)
   model.load_weights(model_path)
-  model.save(os.path.join(hub_destination, "classification"), include_optimizer=False)
+  model.save(
+      os.path.join(hub_destination, "classification"), include_optimizer=False)
 
   # Extracts a sub-model to use pooling feature vector as model output.
   image_input = model.get_layer(index=0).get_output_at(0)
@@ -50,8 +52,7 @@ def export_tfhub(model_path, hub_destination):
 
   # Exports a SavedModel.
   hub_model.save(
-      os.path.join(hub_destination, "feature-vector"),
-      include_optimizer=False)
+      os.path.join(hub_destination, "feature-vector"), include_optimizer=False)
 
 
 def main(argv):