Task delegate 9: Expose BaseOptions in Java AudioClassifier

lu-wang-g · tflite-support-robot · commit 293f2b9f5d54 · 2021-10-01T21:26:20.000-07:00
PiperOrigin-RevId: 400353942
diff --git a/tensorflow_lite_support/cc/task/audio/audio_classifier.h b/tensorflow_lite_support/cc/task/audio/audio_classifier.h
@@ -52,6 +52,13 @@ namespace audio {
 //      `ImageClassifierOptions` used at creation time ("en" by default, i.e.
 //      English). If none of these are available, only the `index` field of the
 //      results will be filled.
+//
+// An example of such model can be found at:
+// https://tfhub.dev/google/lite-model/yamnet/classification/tflite/1
+
+// A CLI demo tool is available for easily trying out this API, and provides
+// example usage. See:
+// https://github.com/tensorflow/tflite-support/tree/master/tensorflow_lite_support/examples/task/audio/desktop
 class AudioClassifier
     : public tflite::task::core::BaseTaskApi<ClassificationResult,
                                              const AudioBuffer&> {
diff --git a/tensorflow_lite_support/java/src/java/org/tensorflow/lite/task/audio/classifier/AudioClassifier.java b/tensorflow_lite_support/java/src/java/org/tensorflow/lite/task/audio/classifier/AudioClassifier.java
@@ -34,6 +34,7 @@
 import org.tensorflow.lite.support.audio.TensorAudio;
 import org.tensorflow.lite.support.audio.TensorAudio.TensorAudioFormat;
 import org.tensorflow.lite.support.tensorbuffer.TensorBuffer;
+import org.tensorflow.lite.task.core.BaseOptions;
 import org.tensorflow.lite.task.core.BaseTaskApi;
 import org.tensorflow.lite.task.core.TaskJniUtils;
 import org.tensorflow.lite.task.core.TaskJniUtils.EmptyHandleProvider;
@@ -64,8 +65,12 @@
  *             use index as label in the result.
  *       </ul>
  * </ul>
+ *
+ * See <a href="https://tfhub.dev/google/lite-model/yamnet/classification/tflite/1">an example</a>
+ * of such model, and <a
+ * href="https://github.com/tensorflow/tflite-support/tree/master/tensorflow_lite_support/examples/task/audio/desktop">a
+ * CLI demo tool</a> for easily trying out this API.
  */
-// TODO(b/182535933): Add a model example and demo comments here.
 public final class AudioClassifier extends BaseTaskApi {
 
   private static final String AUDIO_CLASSIFIER_NATIVE_LIB = "task_audio_jni";
@@ -133,7 +138,11 @@ public long createHandle(
                   long fileDescriptorOffset,
                   AudioClassifierOptions options) {
                 return initJniWithModelFdAndOptions(
-                    fileDescriptor, fileDescriptorLength, fileDescriptorOffset, options);
+                    fileDescriptor,
+                    fileDescriptorLength,
+                    fileDescriptorOffset,
+                    options,
+                    TaskJniUtils.createProtoBaseOptionsHandle(options.getBaseOptions()));
               }
             },
             AUDIO_CLASSIFIER_NATIVE_LIB,
@@ -162,7 +171,8 @@ public long createHandle() {
                       descriptor.getFd(),
                       /*fileDescriptorLength=*/ OPTIONAL_FD_LENGTH,
                       /*fileDescriptorOffset=*/ OPTIONAL_FD_OFFSET,
-                      options);
+                      options,
+                      TaskJniUtils.createProtoBaseOptionsHandle(options.getBaseOptions()));
                 }
               },
               AUDIO_CLASSIFIER_NATIVE_LIB));
@@ -191,7 +201,10 @@ public static AudioClassifier createFromBufferAndOptions(
             new EmptyHandleProvider() {
               @Override
               public long createHandle() {
-                return initJniWithByteBuffer(modelBuffer, options);
+                return initJniWithByteBuffer(
+                    modelBuffer,
+                    options,
+                    TaskJniUtils.createProtoBaseOptionsHandle(options.getBaseOptions()));
               }
             },
             AUDIO_CLASSIFIER_NATIVE_LIB));
@@ -215,6 +228,7 @@ public static class AudioClassifierOptions {
     // 1. java.util.Optional require Java 8 while we need to support Java 7.
     // 2. The Guava library (com.google.common.base.Optional) is avoided in this project. See the
     // comments for labelAllowList.
+    private final BaseOptions baseOptions;
     private final String displayNamesLocale;
     private final int maxResults;
     private final float scoreThreshold;
@@ -233,6 +247,7 @@ public static Builder builder() {
 
     /** A builder that helps to configure an instance of AudioClassifierOptions. */
     public static class Builder {
+      private BaseOptions baseOptions = BaseOptions.builder().build();
       private String displayNamesLocale = "en";
       private int maxResults = -1;
       private float scoreThreshold;
@@ -242,6 +257,12 @@ public static class Builder {
 
       private Builder() {}
 
+      /** Sets the general options to configure Task APIs, such as accelerators. */
+      public Builder setBaseOptions(BaseOptions baseOptions) {
+        this.baseOptions = baseOptions;
+        return this;
+      }
+
       /**
        * Sets the locale to use for display names specified through the TFLite Model Metadata, if
        * any.
@@ -339,13 +360,18 @@ public List<String> getLabelDenyList() {
       return new ArrayList<>(labelDenyList);
     }
 
+    public BaseOptions getBaseOptions() {
+      return baseOptions;
+    }
+
     private AudioClassifierOptions(Builder builder) {
       displayNamesLocale = builder.displayNamesLocale;
       maxResults = builder.maxResults;
       scoreThreshold = builder.scoreThreshold;
       isScoreThresholdSet = builder.isScoreThresholdSet;
       labelAllowList = builder.labelAllowList;
       labelDenyList = builder.labelDenyList;
+      baseOptions = builder.baseOptions;
     }
   }
 
@@ -485,10 +511,11 @@ private static native long initJniWithModelFdAndOptions(
       int fileDescriptor,
       long fileDescriptorLength,
       long fileDescriptorOffset,
-      AudioClassifierOptions options);
+      AudioClassifierOptions options,
+      long baseOptionsHandle);
 
   private static native long initJniWithByteBuffer(
-      ByteBuffer modelBuffer, AudioClassifierOptions options);
+      ByteBuffer modelBuffer, AudioClassifierOptions options, long baseOptionsHandle);
 
   /**
    * Releases memory pointed by {@code nativeHandle}, namely a C++ `AudioClassifier` instance.
diff --git a/tensorflow_lite_support/java/src/native/task/audio/BUILD b/tensorflow_lite_support/java/src/native/task/audio/BUILD
@@ -16,6 +16,7 @@ tflite_jni_binary(
     name = "libtask_audio_jni.so",
     srcs = [
         "//tensorflow_lite_support/java/src/native/task/audio/classifier:audio_classifier_jni.cc",
+        "//tensorflow_lite_support/java/src/native/task/core:task_jni_utils.cc",
     ],
     linkscript = "//tensorflow_lite_support/java:default_version_script.lds",
     deps = [
@@ -25,6 +26,7 @@ tflite_jni_binary(
         "//tensorflow_lite_support/cc/task/audio/proto:audio_classifier_options_cc_proto",
         "//tensorflow_lite_support/cc/task/audio/proto:class_proto_inc",
         "//tensorflow_lite_support/cc/task/audio/proto:classifications_proto_inc",
+        "//tensorflow_lite_support/cc/task/core/proto:base_options_proto_inc",
         "//tensorflow_lite_support/cc/utils:jni_utils",
         "//tensorflow_lite_support/java/jni",
     ],
diff --git a/tensorflow_lite_support/java/src/native/task/audio/classifier/BUILD b/tensorflow_lite_support/java/src/native/task/audio/classifier/BUILD
@@ -18,6 +18,7 @@ tflite_jni_binary(
     name = "libtask_audio_jni.so",
     srcs = [
         "audio_classifier_jni.cc",
+        "//tensorflow_lite_support/java/src/native/task/core:task_jni_utils.cc",
     ],
     linkscript = "//tensorflow_lite_support/java:default_version_script.lds",
     deps = [
@@ -27,6 +28,7 @@ tflite_jni_binary(
         "//tensorflow_lite_support/cc/task/audio/proto:audio_classifier_options_cc_proto",
         "//tensorflow_lite_support/cc/task/audio/proto:class_proto_inc",
         "//tensorflow_lite_support/cc/task/audio/proto:classifications_proto_inc",
+        "//tensorflow_lite_support/cc/task/core/proto:base_options_proto_inc",
         "//tensorflow_lite_support/cc/utils:jni_utils",
         "//tensorflow_lite_support/java/jni",
     ],
diff --git a/tensorflow_lite_support/java/src/native/task/audio/classifier/audio_classifier_jni.cc b/tensorflow_lite_support/java/src/native/task/audio/classifier/audio_classifier_jni.cc
@@ -24,6 +24,7 @@ limitations under the License.
 #include "tensorflow_lite_support/cc/task/audio/proto/audio_classifier_options.pb.h"
 #include "tensorflow_lite_support/cc/task/audio/proto/class_proto_inc.h"
 #include "tensorflow_lite_support/cc/task/audio/proto/classifications_proto_inc.h"
+#include "tensorflow_lite_support/cc/task/core/proto/base_options_proto_inc.h"
 #include "tensorflow_lite_support/cc/utils/jni_utils.h"
 
 namespace {
@@ -38,6 +39,7 @@ using ::tflite::task::audio::AudioClassifier;
 using ::tflite::task::audio::AudioClassifierOptions;
 using ::tflite::task::audio::Class;
 using ::tflite::task::audio::ClassificationResult;
+using ::tflite::task::core::BaseOptions;
 
 // TODO(b/183343074): Share the code below with ImageClassifier.
 
@@ -123,9 +125,16 @@ jobject ConvertToClassificationResults(JNIEnv* env,
 }
 
 // Creates an AudioClassifierOptions proto based on the Java class.
-AudioClassifierOptions ConvertToProtoOptions(JNIEnv* env,
-                                             jobject java_options) {
+AudioClassifierOptions ConvertToProtoOptions(JNIEnv* env, jobject java_options,
+                                             jlong base_options_handle) {
   AudioClassifierOptions proto_options;
+
+  if (base_options_handle != kInvalidPointer) {
+    // proto_options will free the previous base_options and set the new one.
+    proto_options.set_allocated_base_options(
+        reinterpret_cast<BaseOptions*>(base_options_handle));
+  }
+
   jclass java_options_class = env->FindClass(
       "org/tensorflow/lite/task/audio/classifier/"
       "AudioClassifier$AudioClassifierOptions");
@@ -202,9 +211,9 @@ extern "C" JNIEXPORT jlong JNICALL
 Java_org_tensorflow_lite_task_audio_classifier_AudioClassifier_initJniWithModelFdAndOptions(
     JNIEnv* env, jclass thiz, jint file_descriptor,
     jlong file_descriptor_length, jlong file_descriptor_offset,
-    jobject java_options) {
+    jobject java_options, jlong base_options_handle) {
   AudioClassifierOptions proto_options =
-      ConvertToProtoOptions(env, java_options);
+      ConvertToProtoOptions(env, java_options, base_options_handle);
   auto file_descriptor_meta = proto_options.mutable_base_options()
                                   ->mutable_model_file()
                                   ->mutable_file_descriptor_meta();
@@ -220,9 +229,10 @@ Java_org_tensorflow_lite_task_audio_classifier_AudioClassifier_initJniWithModelF
 
 extern "C" JNIEXPORT jlong JNICALL
 Java_org_tensorflow_lite_task_audio_classifier_AudioClassifier_initJniWithByteBuffer(
-    JNIEnv* env, jclass thiz, jobject model_buffer, jobject java_options) {
+    JNIEnv* env, jclass thiz, jobject model_buffer, jobject java_options,
+    jlong base_options_handle) {
   AudioClassifierOptions proto_options =
-      ConvertToProtoOptions(env, java_options);
+      ConvertToProtoOptions(env, java_options, base_options_handle);
   // External proto generated header does not overload `set_file_content` with
   // string_view, therefore GetMappedFileBuffer does not apply here.
   // Creating a std::string will cause one extra copying of data. Thus, the
