Introduce kimm: a collection of image models, blocks and layers written in Keras3

[james77777778]

Repository: https://github.com/james77777778/keras-image-models

If you want to try a different model that isn't available in Keras, you can explore this library.
Feel free to contact me if you encounter any issues or have ideas.

Introduction

Keras Image Models (kimm) aims to provide SOTA models with pretrained weights in a user-friendly manner.

Features:

• 🚀 A model zoo where almost all models come with pre-trained weights on ImageNet.
• 🧰 Providing APIs to export models to .tflite and .onnx.
• 🔧 Supporting the reparameterization technique.
• ✨ Integrated with feature extraction capability.

Usage

import keras
import kimm
import numpy as np


# List available models
print(kimm.list_models("mobileone", weights="imagenet"))  # ['MobileOneS0', ...]

# Initialize model with pretrained ImageNet weights
x = keras.random.uniform([1, 224, 224, 3])
model = kimm.models.MobileOneS0()
y = model.predict(x)
print(y.shape)  # (1, 1000)

# Get reparameterized model by kimm.utils.get_reparameterized_model
reparameterized_model = kimm.utils.get_reparameterized_model(model)
y2 = reparameterized_model.predict(x)
np.testing.assert_allclose(
    keras.ops.convert_to_numpy(y), keras.ops.convert_to_numpy(y2), atol=1e-5
)

# Export model to tflite format
kimm.export.export_tflite(reparameterized_model, 224, "model.tflite")

# Export model to onnx format (note: must be "channels_first" format)
# kimm.export.export_onnx(reparameterized_model, 224, "model.onnx")

# Enable feature extraction by setting `feature_extractor=True`
# `feature_keys` can be optionally specified
model = kimm.models.MobileOneS0(
    include_top=False, feature_extractor=True, feature_keys=["BLOCK2_S16", "BLOCK3_S32"]
)
features = model.predict(x)
for feature_name, feature in features.items():
    print(feature_name, feature.shape)  # BLOCK2_S16 (1, 14, 14, 256) ...

Quickstart

An end-to-end example: fine-tuning an image classification model on a cats vs. dogs dataset

Using kimm.models.EfficientNetLiteB0:

Reference: Transfer learning & fine-tuning (keras.io)

Grad-CAM

Using kimm.models.MobileViTS:

Reference: Grad-CAM class activation visualization (keras.io)

Model Zoo

Model	Paper	Weights are ported from	API
ConvMixer	ICLR 2022 Submission	timm	kimm.models.ConvMixer*
ConvNeXt	CVPR 2022	timm	kimm.models.ConvNeXt*
DenseNet	CVPR 2017	timm	kimm.models.DenseNet*
EfficientNet	ICML 2019	timm	kimm.models.EfficientNet*
EfficientNetLite	ICML 2019	timm	kimm.models.EfficientNetLite*
EfficientNetV2	ICML 2021	timm	kimm.models.EfficientNetV2*
GhostNet	CVPR 2020	timm	kimm.models.GhostNet*
GhostNetV2	NeurIPS 2022	timm	kimm.models.GhostNetV2*
HGNet		timm	kimm.models.HGNet*
HGNetV2		timm	kimm.models.HGNetV2*
InceptionNeXt	arXiv 2023	timm	kimm.models.InceptionNeXt*
InceptionV3	CVPR 2016	timm	kimm.models.InceptionV3
LCNet	arXiv 2021	timm	kimm.models.LCNet*
MobileNetV2	CVPR 2018	timm	kimm.models.MobileNetV2*
MobileNetV3	ICCV 2019	timm	kimm.models.MobileNetV3*
MobileOne	CVPR 2023	timm	kimm.models.MobileOne*
MobileViT	ICLR 2022	timm	kimm.models.MobileViT*
MobileViTV2	arXiv 2022	timm	kimm.models.MobileViTV2*
RegNet	CVPR 2020	timm	kimm.models.RegNet*
RepVGG	CVPR 2021	timm	kimm.models.RepVGG*
ResNet	CVPR 2015	timm	kimm.models.ResNet*
TinyNet	NeurIPS 2020	timm	kimm.models.TinyNet*
VGG	ICLR 2015	timm	kimm.models.VGG*
ViT	ICLR 2021	timm	kimm.models.VisionTransformer*
Xception	CVPR 2017	keras	kimm.models.Xception

[ibinti]

Thanks for your great work.
I am very eager to try out MobileViTV2 with ported timm imagenet weights,
Following MobileOneS0 example, I did

kimm.list_models("mobilevit2", weights="imagenet") 
=> ['MobileViTV2W050', 'MobileViTV2W075', 'MobileViTV2W100', 'MobileViTV2W125', 'MobileViTV2W150', 'MobileViTV2W175', 'MobileViTV2W200']

x = keras.random.uniform([1, 224, 224, 3])
model = kimm.models.MobileViTV2W100()
y = model.predict(x)
print(y.shape)

However, model.predict(x) throws

TypeError: Exception encountered when calling Reshape.call().
cannot reshape array of shape (1, 28, 28, 128) (size 100352) into shape [-1, 16, 2, 16, 2, 128] because the product of specified axis sizes (131072) does not evenly divide 100352
Arguments received by Reshape.call():
  • x=jnp.ndarray(shape=(1, 28, 28, 128), dtype=float32)

Thanks

[ibinti]

Well, I found the reason. MobileViTV2 expects image shape (255,255,3)
x = keras.random.uniform([1, 256, 256, 3])
Thanks

[ibinti]

I wonder how to handle image resizing and scaling. Is there a way to retrieve what the model expects like shape, pixel value range [0,255] float or int , or [0,1] or [-1,1]?

[james77777778]

Hey @ibinti
Thanks for trying out this project.

Actually, there was an update a few days ago to support arbitrary sizes for MobileViT and other ViT models.
https://github.com/james77777778/keras-image-models/releases/tag/0.2.1

I wonder how to handle image resizing and scaling. Is there a way to retrieve what the model expects like shape, pixel value range [0,255] float or int , or [0,1] or [-1,1]?

You can get the default size by:

model = kimm.models.MobileNetV2W050()
print(model._default_size)  # prints 224

You can specify an arbitrary size for the model:

# `input_shape` manner
x = keras.random.uniform([1, 256, 256, 3])
model = kimm.models.MobileNetV2W050(input_shape=[256, 256, 3])
y = model.predict(x)

# `input_tensor` manner
x = keras.random.uniform([1, 256, 256, 3])
inputs = keras.Input([256, 256, 3])
model = kimm.models.MobileNetV2W050(input_tensor=inputs)
y = model.predict(x)

All kimm models expect the value range of [0, 255]. The details of preprocessing can only be inspected in the source code.

• https://github.com/james77777778/keras-image-models/blob/main/kimm/_src/models/mobilevit.py#L571
• https://github.com/james77777778/keras-image-models/blob/main/kimm/_src/models/base_model.py#L147

Maybe I can add this information to the models in the future. Thanks for pointint that out!

[ibinti]

Great. Now I am ready to train kimm.models.MobileViTV2W100 and compare its performance with keras.applications.EfficientNetV2.
Thanks again!