Task delegate 8: Expose BaseOptions in Java ImageSegmenter

PiperOrigin-RevId: 400230011

Author: lu-wang-g

tensorflow_lite_support/java/src/java/org/tensorflow/lite/task/vision/segmenter/ImageSegmenter.java

@@ -30,6 +30,7 @@
 import java.util.List;
 import org.tensorflow.lite.support.image.MlImageAdapter;
 import org.tensorflow.lite.support.image.TensorImage;
+import org.tensorflow.lite.task.core.BaseOptions;
 import org.tensorflow.lite.task.core.TaskJniUtils;
 import org.tensorflow.lite.task.core.TaskJniUtils.EmptyHandleProvider;
 import org.tensorflow.lite.task.core.vision.ImageProcessingOptions;
@@ -186,7 +187,8 @@ public long createHandle() {
                     modelBuffer,
                     options.getDisplayNamesLocale(),
                     options.getOutputType().getValue(),
-                    options.getNumThreads());
+                    TaskJniUtils.createProtoBaseOptionsHandleWithLegacyNumThreads(
+                        options.getBaseOptions(), options.getNumThreads()));
               }
             },
             IMAGE_SEGMENTER_NATIVE_LIB),
@@ -210,6 +212,8 @@ public abstract static class ImageSegmenterOptions {
     private static final OutputType DEFAULT_OUTPUT_TYPE = OutputType.CATEGORY_MASK;
     private static final int NUM_THREADS = -1;
 
+    public abstract BaseOptions getBaseOptions();
+
     public abstract String getDisplayNamesLocale();
 
     public abstract OutputType getOutputType();
@@ -220,13 +224,17 @@ public static Builder builder() {
       return new AutoValue_ImageSegmenter_ImageSegmenterOptions.Builder()
           .setDisplayNamesLocale(DEFAULT_DISPLAY_NAME_LOCALE)
           .setOutputType(DEFAULT_OUTPUT_TYPE)
-          .setNumThreads(NUM_THREADS);
+          .setNumThreads(NUM_THREADS)
+          .setBaseOptions(BaseOptions.builder().build());
     }
 
     /** Builder for {@link ImageSegmenterOptions}. */
     @AutoValue.Builder
     public abstract static class Builder {
 
+      /** Sets the general options to configure Task APIs, such as accelerators. */
+      public abstract Builder setBaseOptions(BaseOptions baseOptions);
+
       /**
        * Sets the locale to use for display names specified through the TFLite Model Metadata, if
        * any.
@@ -245,7 +253,11 @@ public abstract static class Builder {
        *
        * <p>numThreads should be greater than 0 or equal to -1. Setting numThreads to -1 has the
        * effect to let TFLite runtime set the value.
+       *
+       * @deprecated use {@link BaseOptions} to configure number of threads instead. This method
+       *     will override the number of threads configured from {@link BaseOptions}.
        */
+      @Deprecated
       public abstract Builder setNumThreads(int numThreads);
 
       public abstract ImageSegmenterOptions build();
@@ -402,7 +414,8 @@ public long createHandle() {
                     fileDescriptorOffset,
                     options.getDisplayNamesLocale(),
                     options.getOutputType().getValue(),
-                    options.getNumThreads());
+                    TaskJniUtils.createProtoBaseOptionsHandleWithLegacyNumThreads(
+                        options.getBaseOptions(), options.getNumThreads()));
               }
             },
             IMAGE_SEGMENTER_NATIVE_LIB);
@@ -415,10 +428,10 @@ private static native long initJniWithModelFdAndOptions(
       long fileDescriptorOffset,
       String displayNamesLocale,
       int outputType,
-      int numThreads);
+      long baseOptionsHandle);
 
   private static native long initJniWithByteBuffer(
-      ByteBuffer modelBuffer, String displayNamesLocale, int outputType, int numThreads);
+      ByteBuffer modelBuffer, String displayNamesLocale, int outputType, long baseOptionsHandle);
 
   /**
    * The native method to segment the image.

tensorflow_lite_support/java/src/native/task/vision/segmenter/BUILD

@@ -18,6 +18,7 @@ jni_binary_with_tflite(
     name = "libtask_vision_jni.so",
     srcs = [
         "image_segmenter_jni.cc",
+        "//tensorflow_lite_support/java/src/native/task/core:task_jni_utils.cc",
         "//tensorflow_lite_support/java/src/native/task/vision/core:base_vision_task_api_jni.cc",
     ],
     linkscript = "//tensorflow_lite_support/java:default_version_script.lds",
@@ -28,6 +29,7 @@ jni_binary_with_tflite(
     ],
     deps = [
         "//tensorflow_lite_support/cc/port:statusor",
+        "//tensorflow_lite_support/cc/task/core/proto:base_options_proto_inc",
         "//tensorflow_lite_support/cc/task/vision/core:frame_buffer",
         "//tensorflow_lite_support/cc/task/vision/proto:image_segmenter_options_proto_inc",
         "//tensorflow_lite_support/cc/task/vision/proto:segmentations_proto_inc",

tensorflow_lite_support/java/src/native/task/vision/segmenter/image_segmenter_jni.cc

@@ -20,6 +20,7 @@ limitations under the License.
 
 #include "external/com_google_absl/absl/strings/str_cat.h"
 #include "tensorflow_lite_support/cc/port/statusor.h"
+#include "tensorflow_lite_support/cc/task/core/proto/base_options_proto_inc.h"
 #include "tensorflow_lite_support/cc/task/vision/core/frame_buffer.h"
 #include "tensorflow_lite_support/cc/task/vision/image_segmenter.h"
 #include "tensorflow_lite_support/cc/task/vision/proto/image_segmenter_options_proto_inc.h"
@@ -36,6 +37,7 @@ using ::tflite::support::utils::kAssertionError;
 using ::tflite::support::utils::kIllegalArgumentException;
 using ::tflite::support::utils::kInvalidPointer;
 using ::tflite::support::utils::ThrowException;
+using ::tflite::task::core::BaseOptions;
 using ::tflite::task::vision::FrameBuffer;
 using ::tflite::task::vision::ImageSegmenter;
 using ::tflite::task::vision::ImageSegmenterOptions;
@@ -58,9 +60,15 @@ constexpr int kOutputTypeConfidenceMask = 1;
 ImageSegmenterOptions ConvertToProtoOptions(JNIEnv* env,
                                             jstring display_names_locale,
                                             jint output_type,
-                                            jint num_threads) {
+                                            jlong base_options_handle) {
   ImageSegmenterOptions 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));
+  }
+
   const char* pchars = env->GetStringUTFChars(display_names_locale, nullptr);
   proto_options.set_display_names_locale(pchars);
   env->ReleaseStringUTFChars(display_names_locale, pchars);
@@ -78,8 +86,6 @@ ImageSegmenterOptions ConvertToProtoOptions(JNIEnv* env,
                      "Unsupported output type: %d", output_type);
   }
 
-  proto_options.set_num_threads(num_threads);
-
   return proto_options;
 }
 
@@ -185,10 +191,11 @@ extern "C" JNIEXPORT jlong JNICALL
 Java_org_tensorflow_lite_task_vision_segmenter_ImageSegmenter_initJniWithModelFdAndOptions(
     JNIEnv* env, jclass thiz, jint file_descriptor,
     jlong file_descriptor_length, jlong file_descriptor_offset,
-    jstring display_names_locale, jint output_type, jint num_threads) {
+    jstring display_names_locale, jint output_type, jlong base_options_handle) {
   ImageSegmenterOptions proto_options = ConvertToProtoOptions(
-      env, display_names_locale, output_type, num_threads);
-  auto file_descriptor_meta = proto_options.mutable_model_file_with_metadata()
+      env, display_names_locale, output_type, base_options_handle);
+  auto file_descriptor_meta = proto_options.mutable_base_options()
+                                  ->mutable_model_file()
                                   ->mutable_file_descriptor_meta();
   file_descriptor_meta->set_fd(file_descriptor);
   if (file_descriptor_length > 0) {
@@ -203,10 +210,10 @@ Java_org_tensorflow_lite_task_vision_segmenter_ImageSegmenter_initJniWithModelFd
 extern "C" JNIEXPORT jlong JNICALL
 Java_org_tensorflow_lite_task_vision_segmenter_ImageSegmenter_initJniWithByteBuffer(
     JNIEnv* env, jclass thiz, jobject model_buffer,
-    jstring display_names_locale, jint output_type, jint num_threads) {
+    jstring display_names_locale, jint output_type, jlong base_options_handle) {
   ImageSegmenterOptions proto_options = ConvertToProtoOptions(
-      env, display_names_locale, output_type, num_threads);
-  proto_options.mutable_model_file_with_metadata()->set_file_content(
+      env, display_names_locale, output_type, base_options_handle);
+  proto_options.mutable_base_options()->mutable_model_file()->set_file_content(
       static_cast<char*>(env->GetDirectBufferAddress(model_buffer)),
       static_cast<size_t>(env->GetDirectBufferCapacity(model_buffer)));
   return CreateImageSegmenterFromOptions(env, proto_options);