By default, ML Kit's APIs make use of Google trained machine learning models.
These models are designed to cover a wide range of applications. However, some
use cases require models that are more targeted. That is why some ML Kit APIs
now allow you to replace the default models with custom LiteRT models.

Both the [Image Labeling](https://developers.google.com/ml-kit/vision/image-labeling.md.txt) and the
[Object Detection \& Tracking](https://developers.google.com/ml-kit/vision/object-detection.md.txt) API offer support for custom
image classification models. They are compatible with a selection of
high-quality pre-trained models on TensorFlow Hub or your own custom model
trained with TensorFlow or AutoML.

If you need a custom solution for other domains or use-cases, visit the
[On-device Machine Learning page](https://developers.google.com/learn/topics/on-device-ml) for guidance on all of Google's
solutions and tools for on-device machine learning.

## Benefits of using ML Kit with custom models

The benefits for using a custom image classification model with ML Kit are:

- **Easy-to-use high level APIs** - No need to deal with low-level model input/output, handle image pre-/post-processing or building a processing pipeline.
- **No need to worry about label mapping yourself**, ML Kit extracts the labels from LiteRT model metadata and does the mapping for you.
- **Supports custom models from a wide range of sources**, from pre-trained models published on TensorFlow Hub to new models trained with TensorFlow or AutoML.
- Optimized for integration with **Android's Camera APIs**.

And, specifically for [Object Detection and Tracking](https://developers.google.com/ml-kit/vision/object-detection.md.txt):

- **Improve classification accuracy** by locating the objects first and only run the classifier on the related image area.
- **Provide a real-time interactive experience** by providing your users immediate feedback on objects as they are being detected and classified.

## Use a pre-trained image classification model

You can use pre-trained LiteRT models, provided they meet a
[set of criteria](https://developers.google.com/ml-kit/custom-models.md.txt#model-compatibility). Through TensorFlow Hub we are offering
a set of vetted models - from Google or other model creators - that meet these
criteria.

### Use a model published on TensorFlow Hub

[TensorFlow Hub](https://tfhub.dev) offers a wide range of pre-trained image
classification models - from various model creators - that can be used with the
Image Labeling and Object Detection and Tracking APIs. Follow these steps.

1. Pick a model from the [collection of ML Kit compatible models](https://tfhub.dev/ml-kit/collections/image-classification/1).
2. Download the .tflite model file from the model details page. Where available, pick a model format with metadata.
3. Follow our guides for the [Image Labeling API](https://developers.google.com/ml-kit/vision/image-labeling.md.txt#custom-tflite) or [Object Detection and Tracking API](https://developers.google.com/ml-kit/vision/object-detection.md.txt#custom-tflite) on how to bundle model file with your project and use it in your Android or iOS application.

## Train your own image classification model

If no pre-trained image classification model fits your needs, there are various
ways to train your own LiteRT model, some of which are outlined and
discussed in more detail in the following section.

| Options to train your own image classification model ||
|---|---|
| **AutoML** | - Offered through Google Cloud AI - Create state-of the art image classification models - Evaluate between performance and size |
| **Convert a TensorFlow model to LiteRT** | - Train a model with TensorFlow and then convert it to LiteRT |

### AutoML

Image classification models trained using
[AutoML](https://docs.cloud.google.com/gemini-enterprise-agent-platform/machine-learning/beginner/beginners-guide) are supported by the
custom models in the [Image Labeling](https://developers.google.com/ml-kit/vision/image-labeling.md.txt) and
[Object Detection and Tracking API](https://developers.google.com/ml-kit/vision/object-detection.md.txt) APIs. These APIs also
support download of models that are hosted with [Cloud Storage](https://firebase.google.com/docs/storage).

To learn more about how to use a model trained with AutoML
in your Android and iOS apps, follow the custom model guides for each
API, depending on your use case.

> [!NOTE]
> **Note:** ML Kit only supports custom image classification models. Although AutoML allows training of object detection models, these cannot be used with ML Kit.

### Models created using the LiteRT converter

If you have an existing TensorFlow image classification model, you can convert
it using the [LiteRT converter](https://www.tensorflow.org/lite/convert). Ensure the model created
meets the following compatibility requirements.

To learn more about how to use a LiteRT model in your Android and iOS apps,
follow our guides for the [Image Labeling API](https://developers.google.com/ml-kit/vision/image-labeling.md.txt) or the
[Object Detection and Tracking API](https://developers.google.com/ml-kit/vision/object-detection.md.txt), depending on your use
case.

## LiteRT model compatibility

You can use any pre-trained LiteRT image classification model, provided it meets
these requirements:

### Tensors

- The model must have only one input tensor with the following constraints:
  - The data is in RGB pixel format.
  - The data is UINT8 or FLOAT32 type. If the input tensor type is FLOAT32, it must specify the NormalizationOptions by attaching [Metadata](https://developers.google.com/ml-kit/custom-models.md.txt#metadata).
  - The tensor has 4 dimensions : BxHxWxC, where:
    - B is the batch size. It must be 1 (inference on larger batches is not supported).
    - W and H are the input width and height.
    - C is the number of expected channels. It must be 3.
- The model must have at least one output tensor with N classes and either 2 or 4 dimensions:
  - (1xN)
  - (1x1x1xN)
- Only single-head models are fully supported. Multi-head models may output unexpected results.

### Metadata

You can add metadata to the LiteRT file as explained in
[Adding metadata to LiteRT models](https://www.tensorflow.org/lite/convert/metadata).

To use a model with FLOAT32 input tensor, you must specify the
[NormalizationOptions](https://github.com/tensorflow/tensorflow/blob/d64c20eb62c5cfb9ff5936204afb8fb7c83cfc84/tensorflow/lite/experimental/support/metadata/metadata_schema.fbs#L295-L318) in the metadata.

We also recommend that you attach this metadata to the output tensor
[TensorMetadata](https://github.com/tensorflow/tensorflow/blob/d64c20eb62c5cfb9ff5936204afb8fb7c83cfc84/tensorflow/lite/experimental/support/metadata/metadata_schema.fbs#L396-L397):

- A label map specifying the name of each output class, as an [AssociatedFile](https://github.com/tensorflow/tensorflow/blob/d64c20eb62c5cfb9ff5936204afb8fb7c83cfc84/tensorflow/lite/experimental/support/metadata/metadata_schema.fbs#L434-L438) with type [TENSOR_AXIS_LABELS](https://github.com/tensorflow/tensorflow/blob/d64c20eb62c5cfb9ff5936204afb8fb7c83cfc84/tensorflow/lite/experimental/support/metadata/metadata_schema.fbs#L42-L54) (otherwise only the numerical output class indexes can be returned)
- A default score threshold under which results are considered too low-confidence to be returned, as a [ProcessUnit](https://github.com/tensorflow/tensorflow/blob/d64c20eb62c5cfb9ff5936204afb8fb7c83cfc84/tensorflow/lite/experimental/support/metadata/metadata_schema.fbs#L422-L429) with [ScoreThresholdingOptions](https://github.com/tensorflow/tensorflow/blob/d64c20eb62c5cfb9ff5936204afb8fb7c83cfc84/tensorflow/lite/experimental/support/metadata/metadata_schema.fbs#L360-L366)