Volo D2 save with preprocessing layers

[HolgerMoritz]

Trying to save a Volo model runs into the exception 'class Ellipsis ... has no get_config', the background is, that I wanted to make this model a typical 'tensorflow-style' model with preprocessing layer & softmax output included. Here's my code (tensorflow 2.18.0):

from keras_cv_attention_models import model_surgery,volo
import tensorflow as tf
from skimage.data import chelsea
from skimage import io
from tensorflow.keras.preprocessing import image
import numpy as np

def get_preprocessed_image_from_file(file_path, image_size, normalize):
img = image.load_img(file_path)
x = image.img_to_array(img)
x = tf.image.resize(x,(image_size, image_size))
if normalize == True:
mean = np.array([123.68, 116.78, 103.94]).astype("float32") #mean = np.array([0.485, 0.456, 0.406]).astype("float32") * 255.0
std = np.array([58.393, 57.12, 57.375]).astype("float32")#std = np.array([0.229, 0.224, 0.225]).astype("float32") * 255.0
x = (x - mean) / std
print("*** normalized ***")

x = np.expand_dims(x, axis=0)
return x

mm = volo.VOLO_d2(pretrained="models/volo_d2_384_imagenet.h5")

img = io.imread('images/chelsea.bmp') # bmp is loss-free! jpg is compressed with losses!, so we can better compare inference results
img = tf.keras.applications.imagenet_utils.preprocess_input(img, mode='torch') # Chelsea the cat (torch)

pred = mm(tf.expand_dims(tf.image.resize(img, mm.input_shape[1:3]), 0)).numpy()
pred = tf.nn.softmax(pred).numpy() # If classifier activation is not softmax
print(tf.keras.applications.imagenet_utils.decode_predictions(pred)[0])
print("----------------")

imm = tf.keras.applications.imagenet_utils.preprocess_input(chelsea()) # Chelsea the cat (legacy)

#from legacy
mean = [103.939, 116.779, 123.68]
imm += mean
imm = imm[..., ::-1]

to torch

imm /= 255.
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
imm -= mean
imm /= std

pred = mm(tf.expand_dims(tf.image.resize(imm, mm.input_shape[1:3]), 0)).numpy()
pred = tf.nn.softmax(pred).numpy() # If classifier activation is not softmax
print(tf.keras.applications.imagenet_utils.decode_predictions(pred)[0])
print("----------------")

save model including preprocessing (NOT yet working, exception)

==================================

mean = [123.68, 116.78, 103.94]
std = [58.393, 57.12, 57.375]
variance = [s**2 for s in std]

inputs = tf.keras.Input(shape=(384, 384, 3)) # Adjust the shape as necessary
preprocessed = tf.keras.layers.Normalization(mean=mean, variance=variance)(inputs)

mm = volo.VOLO_d2(input_shape=(384, 384, 3),pretrained="models/volo_d2_384_imagenet.h5")
pred = mm(preprocessed) # Use the existing model to get outputs
outputs= tf.nn.softmax(pred)

Create the new model & save

new_model = tf.keras.Model(inputs=inputs, outputs=outputs)
new_model._name = 'volo_d2_384'
new_model.compile(optimizer='adam', loss=tf.keras.losses.CategoricalFocalCrossentropy(), metrics=['accuracy'])
new_model.save('volo_d2_384.h5') # exception
model_surgery.export_onnx(new_model, fuse_conv_bn=True, batch_size=1, simplify=False)

[HolgerMoritz]

Additional Remark:
Using tf.keras.Model with save does work without issue with several other model architectures from your collection.

Any help is appreciated