Android ImageSegmenter does not parse categories from the model metadata

[presidentio]

Have I written custom code (as opposed to using a stock example script provided in MediaPipe)

Yes

OS Platform and Distribution

Android 16

Mobile device if the issue happens on mobile device

Pixel 8a

Browser and version if the issue happens on browser

No response

Programming Language and version

Kotlin

MediaPipe version

0.10.29, 0.10.21

Bazel version

No response

Solution

Multy Class Segmentation

Android Studio, NDK, SDK versions (if issue is related to building in Android environment)

Android Studio 2025.1.2, SDK - compile: 36, targetSdk:35

Xcode & Tulsi version (if issue is related to building for iOS)

No response

Describe the actual behavior

ImageSegmenter.segment returns category mask with category ids 253, 254, 255

Describe the expected behaviour

ImageSegmenter.segment returns category mask with category ids 0, 1, 2

Standalone code/steps you may have used to try to get what you need

https://drive.google.com/file/d/1vBlzjtaxCl00qWNK70yTd256z3dcx5Ot/view?usp=sharing

Other info / Complete Logs

I've trained a custom model with 3 categories https://drive.google.com/file/d/1vBlzjtaxCl00qWNK70yTd256z3dcx5Ot/view?usp=sharing.
Code to reproduce the issue:

val options = ImageSegmenterOptions.builder()
            .setBaseOptions(
                BaseOptions.builder().setModelAssetPath(modelPath).build()
            )
            .setRunningMode(RunningMode.IMAGE)
            .setOutputCategoryMask(true)
            .setOutputConfidenceMasks(false)
            .build();
val segmenter = ImageSegmenter.createFromOptions(context, options)
val mpImage = BitmapImageBuilder(image).build()
val segmentationResult = segmenter.segment(mpImage)
val mpMask = segmentationResult.categoryMask().get()
val maskByteArray = ByteBufferExtractor.extract(mpMask)
val set = mutableSetOf<Any>()
for (i in 0 until maskByteArray.capacity()) {
   set += maskByteArray.get(i) and 0xFF
}
Log.d("vision", "Unique values: $set")

ImageSegmenter on python with the same model returns categories 0, 1, 2.
Also the same model works in mediapipe studio https://mediapipe-studio.webapps.google.com/studio/demo/image_segmenter