Make hub.ImageModuleInfo, hub.attach_image_module_info() visible

for use by module publishers. PiperOrigin-RevId: 212239584
yangxiao5201 · Sep 14, 2018 · d3ca254 · d3ca254
1 parent e1eabdd
commit d3ca254
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 8 deletions.
diff --git a/tensorflow_hub/BUILD b/tensorflow_hub/BUILD
@@ -123,6 +123,7 @@ py_library(
     srcs_version = "PY2AND3",
     deps = [
         ":all_protos_py_pb2",
+        ":native_module",
     ],
 )
 

diff --git a/tensorflow_hub/__init__.py b/tensorflow_hub/__init__.py
@@ -29,8 +29,10 @@
 from tensorflow_hub.estimator import register_module_for_export
 from tensorflow_hub.feature_column import image_embedding_column
 from tensorflow_hub.feature_column import text_embedding_column
+from tensorflow_hub.image_util import attach_image_module_info
 from tensorflow_hub.image_util import get_expected_image_size
 from tensorflow_hub.image_util import get_num_image_channels
+from tensorflow_hub.image_util import ImageModuleInfo
 from tensorflow_hub.module import load_module_spec
 from tensorflow_hub.module import Module
 from tensorflow_hub.module_spec import ModuleSpec
@@ -53,8 +55,10 @@
     "register_module_for_export",
     "image_embedding_column",
     "text_embedding_column",
+    "attach_image_module_info",
     "get_expected_image_size",
     "get_num_image_channels",
+    "ImageModuleInfo",
     "Module",
     "ModuleSpec",
     "add_signature",

diff --git a/tensorflow_hub/image_util.py b/tensorflow_hub/image_util.py
@@ -19,17 +19,27 @@
 from __future__ import print_function
 
 from tensorflow_hub import image_module_info_pb2
+from tensorflow_hub import native_module
 
 
-# Image modules can provide further information for the utilities in this file
-# by attaching an ImageModuleInfo message under this key.
+# hub.Modules for images can provide further information for the utilities
+# in this file by attaching an ImageModuleInfo message under this key.
 IMAGE_MODULE_INFO_KEY = "image_module_info"
 
 
-def get_image_module_info(module_or_spec):
+# The externally visible name of the message is hub.ImageModuleInfo
+ImageModuleInfo = image_module_info_pb2.ImageModuleInfo  # pylint: disable=invalid-name
+
+
+def attach_image_module_info(image_module_info):
+  """Attaches an ImageModuleInfo message from within a module_fn."""
+  native_module.attach_message(IMAGE_MODULE_INFO_KEY, image_module_info)
+
+
+def get_image_module_info(module_or_spec, required=False):
   """Returns the module's attached ImageModuleInfo message, or None."""
   return module_or_spec.get_attached_message(
-      IMAGE_MODULE_INFO_KEY, image_module_info_pb2.ImageModuleInfo)
+      IMAGE_MODULE_INFO_KEY, ImageModuleInfo, required=required)
 
 
 def get_expected_image_size(module_or_spec, signature=None, input_name=None):

diff --git a/tensorflow_hub/image_util_test.py b/tensorflow_hub/image_util_test.py
@@ -19,7 +19,6 @@
 from __future__ import print_function
 
 import tensorflow as tf
-from tensorflow_hub import image_module_info_pb2
 from tensorflow_hub import image_util
 from tensorflow_hub import module
 from tensorflow_hub import native_module
@@ -39,11 +38,10 @@ def image_module_fn_with_info():
   sum_all = tf.reduce_sum(images, [1, 2, 3])
   native_module.add_signature(inputs=dict(images=images),
                               outputs=dict(default=sum_all))
-  image_module_info = image_module_info_pb2.ImageModuleInfo()
+  image_module_info = image_util.ImageModuleInfo()
   size = image_module_info.default_image_size
   size.height, size.width = 2, 4
-  native_module.attach_message(image_util.IMAGE_MODULE_INFO_KEY,
-                               image_module_info)
+  image_util.attach_image_module_info(image_module_info)
 
 
 class ImageModuleTest(tf.test.TestCase):