Improved API docs for TFLite Task Lib

PiperOrigin-RevId: 447372882

Author: khanhlvg

tensorflow_lite_support/opensource/opensource_only.files

@@ -44,6 +44,7 @@ tensorflow_lite_support/tools/pip_package/build_pip_package.sh:
 tensorflow_lite_support/tools/pip_package/metadata_writers.__init__.py:
 tensorflow_lite_support/tools/pip_package/setup.py:
 tensorflow_lite_support/tools/pip_package/simple_console_for_windows.py:
+tensorflow_lite_support/tools/pip_package/task.__init__.py:
 tensorflow_lite_support/tools/pip_package/task_audio.__init__.py:
 tensorflow_lite_support/tools/pip_package/task_core.__init__.py:
 tensorflow_lite_support/tools/pip_package/task_processor.__init__.py:

tensorflow_lite_support/python/task/audio/core/tensor_audio.py

@@ -18,14 +18,14 @@
 from tensorflow_lite_support.python.task.audio.core import audio_record
 from tensorflow_lite_support.python.task.audio.core.pybinds import _pywrap_audio_buffer
 
-_CppAudioFormat = _pywrap_audio_buffer.AudioFormat
 _LoadAudioBufferFromFile = _pywrap_audio_buffer.LoadAudioBufferFromFile
+AudioFormat = _pywrap_audio_buffer.AudioFormat
 
 
 class TensorAudio(object):
   """A wrapper class to store the input audio."""
 
-  def __init__(self, audio_format: _CppAudioFormat, buffer_size: int) -> None:
+  def __init__(self, audio_format: AudioFormat, buffer_size: int) -> None:
     """Initializes the `TensorAudio` object.
 
     Args:
@@ -137,7 +137,7 @@ def load_from_array(self,
       self._buffer[-shift:, :] = src[offset:offset + size].copy()
 
   @property
-  def format(self) -> _CppAudioFormat:
+  def format(self) -> AudioFormat:
     """Gets the audio format of the audio."""
     return self._format
 

tensorflow_lite_support/tools/pip_package/build_pip_package.sh

@@ -98,6 +98,7 @@ function prepare_src() {
     # Task Library is not supported on Windows yet.
     mkdir ${TMPDIR}/tflite_support/task
     mkdir ${TMPDIR}/tflite_support/task/core
+    cp tensorflow_lite_support/tools/pip_package/task.__init__.py ${TMPDIR}/tflite_support/task/__init__.py
     cp tensorflow_lite_support/tools/pip_package/task_core.__init__.py ${TMPDIR}/tflite_support/task/core/__init__.py
     mkdir ${TMPDIR}/tflite_support/task/vision
     cp tensorflow_lite_support/tools/pip_package/task_vision.__init__.py ${TMPDIR}/tflite_support/task/vision/__init__.py

tensorflow_lite_support/tools/pip_package/task.__init__.py

@@ -0,0 +1,32 @@
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+# ==============================================================================
+"""The TensorFlow Lite Task Library.
+
+TensorFlow Lite Task Library contains a set of powerful and easy-to-use
+task-specific libraries for app developers to create ML experiences with
+TensorFlow Lite. It provides optimized out-of-box model interfaces for popular
+machine learning tasks, such as image and text classification. The model
+interfaces are specifically designed for each task to achieve the best
+performance and usability.
+
+Read more in the [Task Library Guide](
+https://tensorflow.org/lite/inference_with_metadata/task_library/overview).
+"""
+
+from . import audio
+from . import core
+from . import processor
+from . import text
+from . import vision

tensorflow_lite_support/tools/pip_package/task_audio.__init__.py

@@ -12,7 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""An import entry for Task Audio Library."""
+"""TensorFlow Lite Task Library Audio APIs.
+
+This module provides interface to run TensorFlow Lite audio models.
+"""
 
 from tensorflow_lite_support.python.task.audio import audio_classifier
 from tensorflow_lite_support.python.task.audio import audio_embedder
@@ -24,4 +27,11 @@
 AudioEmbedder = audio_embedder.AudioEmbedder
 AudioEmbedderOptions = audio_embedder.AudioEmbedderOptions
 AudioRecord = audio_record.AudioRecord
+AudioFormat = tensor_audio.AudioFormat
 TensorAudio = tensor_audio.TensorAudio
+
+# Remove unnecessary modules to avoid duplication in API docs.
+del audio_classifier
+del audio_embedder
+del audio_record
+del tensor_audio

tensorflow_lite_support/tools/pip_package/task_core.__init__.py

@@ -12,8 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""An import entry for the Task Core module."""
+"""TensorFlow Lite Task Library's core module.
+
+This module contains classes used across multiple tasks in the Task Library."""
 
 from tensorflow_lite_support.python.task.core.proto import base_options_pb2
 
 BaseOptions = base_options_pb2.BaseOptions
+
+# Remove unnecessary modules to avoid duplication in API docs.
+del base_options_pb2

tensorflow_lite_support/tools/pip_package/task_processor.__init__.py

@@ -12,7 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""An import entry for the Task Processor module."""
+"""TensorFlow Lite Task Library's processor module.
+
+This module contains classes related to the pre-processing and post-processing
+steps of the Task Library.
+"""
 
 from tensorflow_lite_support.python.task.processor.proto import bounding_box_pb2
 from tensorflow_lite_support.python.task.processor.proto import classification_options_pb2
@@ -37,3 +41,16 @@
 Segmentation = segmentations_pb2.Segmentation
 SearchOptions = search_options_pb2.SearchOptions
 SearchResult = search_result_pb2.SearchResult
+
+# Remove unnecessary modules to avoid duplication in API docs.
+del bounding_box_pb2
+del classification_options_pb2
+del classifications_pb2
+del detection_options_pb2
+del detections_pb2
+del embedding_options_pb2
+del embedding_pb2
+del segmentation_options_pb2
+del segmentations_pb2
+del search_options_pb2
+del search_result_pb2

tensorflow_lite_support/tools/pip_package/task_text.__init__.py

@@ -12,7 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""An import entry for Task Text Library."""
+"""TensorFlow Lite Task Library Text APIs.
+
+This module provides interface to run TensorFlow Lite natural language
+processing models.
+"""
 
 from tensorflow_lite_support.python.task.text import text_embedder
 from tensorflow_lite_support.python.task.text import text_searcher
@@ -21,3 +25,7 @@
 TextEmbedderOptions = text_embedder.TextEmbedderOptions
 TextSearcher = text_searcher.TextSearcher
 TextSearcherOptions = text_searcher.TextSearcherOptions
+
+# Remove unnecessary modules to avoid duplication in API docs.
+del text_embedder
+del text_searcher

tensorflow_lite_support/tools/pip_package/task_vision.__init__.py

@@ -12,7 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""An import entry for Task Vision Library."""
+"""TensorFlow Lite Task Library Vision APIs.
+
+This module provides interface to run TensorFlow Lite computer vision models.
+"""
 
 from tensorflow_lite_support.python.task.vision import image_classifier
 from tensorflow_lite_support.python.task.vision import image_embedder
@@ -32,3 +35,11 @@
 ImageSearcher = image_searcher.ImageSearcher
 ImageSearcherOptions = image_searcher.ImageSearcherOptions
 TensorImage = tensor_image.TensorImage
+
+# Remove unnecessary modules to avoid duplication in API docs.
+del image_classifier
+del image_embedder
+del image_segmenter
+del image_searcher
+del object_detector
+del tensor_image

tensorflow_lite_support/tools/pip_package/tflite_support.__init__.py

@@ -12,17 +12,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""The TFLite Support Library.
+"""The TensorFlow Lite Support Library.
 
 Install the pip package:
 
 ```
 pip install tflite-support
 ```
 
+This package provides two major features:
+* Metadata writers: add metadata to TensorFlow Lite models.
+* Task Library: run TensorFlow Lite models of major machine learning tasks.
+
 To learn more about metadata, flatbuffers and TensorFlow Lite models, check out
 the [metadata section](https://www.tensorflow.org/lite/convert/metadata) of the
-TF Lite guide.
+TensorFlow Lite guide.
+
+To learn more about Task Library, check out the
+[documentation](https://www.tensorflow.org/lite/inference_with_metadata/task_library/overview)
+on the TensorFlow Lite website.
 """
 
 # In the original project structure, all python targets are accessed by paths