Update app.py

app.py

@@ -3,11 +3,28 @@ import tensorflow as tf
 from PIL import Image
 import numpy as np
 
-# Cargamos el modelo guardado
+from transformers import TFDeiTModel
+
+backbone = TFDeiTModel.from_pretrained("facebook/deit-tiny-patch16-224")
+
+def transpose_channels(x):
+  return tf.transpose(x, [0, 3, 1, 2])
+
+def deit_forward(x):
+  return backbone(pixel_values=x).last_hidden_state[:, 0, :]
+
 tf.keras.config.enable_unsafe_deserialization()
-model = tf.keras.models.load_model("model.keras", safe_mode=False)
 
-# Tus clases
+model = tf.keras.models.load_model(
+  "model.keras",
+  custom_objects={
+      "transpose_channels": transpose_channels,
+      "deit_forward": deit_forward
+  },
+  safe_mode=False
+)
+
+# Clases
 labels = ['pituitary', 'meningioma', 'notumor', 'glioma']
 
 # Preprocesamiento (ajustado al modelo Deit-Tiny 224x224)
@@ -33,4 +50,4 @@ interface = gr.Interface(
     description="Clasificador entrenado en 4 categorías: pituitary, meningioma, notumor, glioma"
 )
 
-interface.launch()
+interface.launch()