Throw RuntimeExceptions instead of AssertionError in ImageSegmenter

PiperOrigin-RevId: 400353917

Author: lu-wang-g

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

@@ -87,8 +87,9 @@ public final class ImageSegmenter extends BaseVisionTaskApi {
    *
    * @param modelPath path of the segmentation model with metadata in the assets
    * @throws IOException if an I/O error occurs when loading the tflite model
-   * @throws AssertionError if error occurs when creating {@link ImageSegmenter} from the native
-   *     code
+   * @throws IllegalArgumentException if an argument is invalid
+   * @throws IllegalStateException if there is an internal error
+   * @throws RuntimeException if there is an otherwise unspecified error
    */
   public static ImageSegmenter createFromFile(Context context, String modelPath)
       throws IOException {
@@ -100,8 +101,9 @@ public static ImageSegmenter createFromFile(Context context, String modelPath)
    *
    * @param modelFile the segmentation model {@link File} instance
    * @throws IOException if an I/O error occurs when loading the tflite model
-   * @throws AssertionError if error occurs when creating {@link ImageSegmenter} from the native
-   *     code
+   * @throws IllegalArgumentException if an argument is invalid
+   * @throws IllegalStateException if there is an internal error
+   * @throws RuntimeException if there is an otherwise unspecified error
    */
   public static ImageSegmenter createFromFile(File modelFile) throws IOException {
     return createFromFileAndOptions(modelFile, ImageSegmenterOptions.builder().build());
@@ -113,8 +115,8 @@ public static ImageSegmenter createFromFile(File modelFile) throws IOException {
    *
    * @param modelBuffer a direct {@link ByteBuffer} or a {@link MappedByteBuffer} of the
    *     segmentation model
-   * @throws AssertionError if error occurs when creating {@link ImageSegmenter} from the native
-   *     code
+   * @throws IllegalStateException if there is an internal error
+   * @throws RuntimeException if there is an otherwise unspecified error
    * @throws IllegalArgumentException if the model buffer is not a direct {@link ByteBuffer} or a
    *     {@link MappedByteBuffer}
    */
@@ -127,8 +129,9 @@ public static ImageSegmenter createFromBuffer(final ByteBuffer modelBuffer) {
    *
    * @param modelPath path of the segmentation model with metadata in the assets
    * @throws IOException if an I/O error occurs when loading the tflite model
-   * @throws AssertionError if error occurs when creating {@link ImageSegmenter} from the native
-   *     code
+   * @throws IllegalArgumentException if an argument is invalid
+   * @throws IllegalStateException if there is an internal error
+   * @throws RuntimeException if there is an otherwise unspecified error
    */
   public static ImageSegmenter createFromFileAndOptions(
       Context context, String modelPath, final ImageSegmenterOptions options) throws IOException {
@@ -146,8 +149,9 @@ public static ImageSegmenter createFromFileAndOptions(
    *
    * @param modelFile the segmentation model {@link File} instance
    * @throws IOException if an I/O error occurs when loading the tflite model
-   * @throws AssertionError if error occurs when creating {@link ImageSegmenter} from the native
-   *     code
+   * @throws IllegalArgumentException if an argument is invalid
+   * @throws IllegalStateException if there is an internal error
+   * @throws RuntimeException if there is an otherwise unspecified error
    */
   public static ImageSegmenter createFromFileAndOptions(
       File modelFile, final ImageSegmenterOptions options) throws IOException {
@@ -167,8 +171,8 @@ public static ImageSegmenter createFromFileAndOptions(
    *
    * @param modelBuffer a direct {@link ByteBuffer} or a {@link MappedByteBuffer} of the
    *     segmentation model
-   * @throws AssertionError if error occurs when creating {@link ImageSegmenter} from the native
-   *     code
+   * @throws IllegalStateException if there is an internal error
+   * @throws RuntimeException if there is an otherwise unspecified error
    * @throws IllegalArgumentException if the model buffer is not a direct {@link ByteBuffer} or a
    *     {@link MappedByteBuffer}
    */
@@ -282,7 +286,8 @@ public abstract static class Builder {
    *     Segmentation} element is expected to be returned. The result is stored in a {@link List}
    *     for later extension to e.g. instance segmentation models, which may return one segmentation
    *     per object.
-   * @throws AssertionError if error occurs when segmenting the image from the native code
+   * @throws IllegalStateException if there is an internal error
+   * @throws RuntimeException if there is an otherwise unspecified error
    * @throws IllegalArgumentException if the color space type of image is unsupported
    */
   public List<Segmentation> segment(TensorImage image) {
@@ -315,7 +320,8 @@ public List<Segmentation> segment(TensorImage image) {
    *     Segmentation} element is expected to be returned. The result is stored in a {@link List}
    *     for later extension to e.g. instance segmentation models, which may return one segmentation
    *     per object.
-   * @throws AssertionError if error occurs when segmenting the image from the native code
+   * @throws IllegalStateException if there is an internal error
+   * @throws RuntimeException if there is an otherwise unspecified error
    * @throws IllegalArgumentException if the color space type of image is unsupported
    */
   public List<Segmentation> segment(TensorImage image, ImageProcessingOptions options) {
@@ -339,7 +345,8 @@ public List<Segmentation> run(
    *     Segmentation} element is expected to be returned. The result is stored in a {@link List}
    *     for later extension to e.g. instance segmentation models, which may return one segmentation
    *     per object.
-   * @throws AssertionError if error occurs when segmenting the image from the native code
+   * @throws IllegalStateException if there is an internal error
+   * @throws RuntimeException if there is an otherwise unspecified error
    * @throws IllegalArgumentException if the storage type or format of the image is unsupported
    */
   public List<Segmentation> segment(MlImage image) {
@@ -364,7 +371,8 @@ public List<Segmentation> segment(MlImage image) {
    *     Segmentation} element is expected to be returned. The result is stored in a {@link List}
    *     for later extension to e.g. instance segmentation models, which may return one segmentation
    *     per object.
-   * @throws AssertionError if error occurs when segmenting the image from the native code
+   * @throws IllegalStateException if there is an internal error
+   * @throws RuntimeException if there is an otherwise unspecified error
    * @throws IllegalArgumentException if the color space type of image is unsupported
    */
   public List<Segmentation> segment(MlImage image, ImageProcessingOptions options) {

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

@@ -33,8 +33,9 @@ namespace {
 
 using ::tflite::support::StatusOr;
 using ::tflite::support::utils::CreateByteArray;
-using ::tflite::support::utils::kAssertionError;
+using ::tflite::support::utils::GetExceptionClassNameForStatusCode;
 using ::tflite::support::utils::kIllegalArgumentException;
+using ::tflite::support::utils::kIllegalStateException;
 using ::tflite::support::utils::kInvalidPointer;
 using ::tflite::support::utils::ThrowException;
 using ::tflite::task::core::BaseOptions;
@@ -97,7 +98,7 @@ void ConvertFromSegmentationResults(JNIEnv* env,
   if (results.segmentation_size() != 1) {
     // Should never happen.
     ThrowException(
-        env, kAssertionError,
+        env, kIllegalStateException,
         "ImageSegmenter only supports one segmentation result, getting %d",
         results.segmentation_size());
   }
@@ -171,13 +172,17 @@ jlong CreateImageSegmenterFromOptions(JNIEnv* env,
   if (image_segmenter_or.ok()) {
     return reinterpret_cast<jlong>(image_segmenter_or->release());
   } else {
-    ThrowException(env, kAssertionError,
-                   "Error occurred when initializing ImageSegmenter: %s",
-                   image_segmenter_or.status().message().data());
+    ThrowException(
+        env,
+        GetExceptionClassNameForStatusCode(image_segmenter_or.status().code()),
+        "Error occurred when initializing ImageSegmenter: %s",
+        image_segmenter_or.status().message().data());
     return kInvalidPointer;
   }
 }
 
+}  // namespace
+
 extern "C" JNIEXPORT void JNICALL
 Java_org_tensorflow_lite_task_vision_segmenter_ImageSegmenter_deinitJni(
     JNIEnv* env, jobject thiz, jlong native_handle) {
@@ -233,10 +238,9 @@ Java_org_tensorflow_lite_task_vision_segmenter_ImageSegmenter_segmentNative(
     ConvertFromSegmentationResults(env, results_or.value(), jmask_buffers,
                                    jmask_shape, jcolored_labels);
   } else {
-    ThrowException(env, kAssertionError,
-                   "Error occurred when segmenting the image: %s",
-                   results_or.status().message().data());
+    ThrowException(
+        env, GetExceptionClassNameForStatusCode(results_or.status().code()),
+        "Error occurred when segmenting the image: %s",
+        results_or.status().message().data());
   }
 }
-
-}  // namespace