import keras
img_height = 162
img_width = 300
num_classes = 5

# base model for transfer learning
base_model = keras.applications.DenseNet121(
    input_shape=(img_height, img_width, 3),
    include_top=False,
)
base_model.trainable = False  # Freeze the base model

model = keras.models.Sequential(
    [
        keras.layers.Input((img_height, img_width, 1)),
        keras.layers.Lambda(
            lambda x: tf.repeat(
                x,
                3,
                axis=3,
            )
        ),  # Convert grayscale to RGB
        keras.layers.Lambda(keras.applications.densenet.preprocess_input),
        base_model,
        keras.layers.GlobalAveragePooling2D(),
        keras.layers.BatchNormalization(),
        keras.layers.Dense(256, activation="relu"),
        keras.layers.Dropout(0.5),
        keras.layers.Dense(num_classes, activation="softmax"),
    ]
)
# Load the trained weights
model.load_weights('hf://c2p-cmd/knee_oa_classifier')