Howto use Mediapipe to create MLKit object detection model

[cjohn001]

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

None

OS Platform and Distribution

Ubuntu 24.4

Python Version

3.11

MediaPipe Model Maker version

No response

Task name (e.g. Image classification, Gesture recognition etc.)

Object detection

Describe the actual behavior

Cannot execute model with MLKIT

Describe the expected behaviour

can execute model with MLKIT

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

spec = object_detector.SupportedModels.MOBILENET_V2_I320
hparams = object_detector.HParams(export_dir='exported_model')
options = object_detector.ObjectDetectorOptions(
    supported_model=spec,
    hparams=hparams
)
model = object_detector.ObjectDetector.create(
    train_data=train_data,
    validation_data=validation_data,
    options=options)
model.export_model()

Other info / Complete Logs

Hello, I am currently trying to create a custom object detection model with Mediapipe model maker in order to use it with MLKit an Android and iOS devices. However, it seems the output shape of the generated model seems to not match what MLKit is expecting. I get bouding boxes
but when enabling classification  

customObjectDetectionOptions.shouldEnableClassification = true

I see the following error message:

failed to detect object with error Pipeline failed to fully start:
CalculatorGraph::Run() failed: 
Calculator::Open() for node "BoxClassifierCalculator" failed: #vk Unexpected number of dimensions for output index 0: got 3D, expected either 2D (BxN with B=1) or 4D (BxHxWxN with B=1, W=1, H=1)..

Is there a way to adopt the models output shapes in order to get them to work with mlkit?