Upload inference.py

inference.py

@@ -0,0 +1,46 @@
+import tensorflow as tf
+import numpy as np
+from PIL import Image
+import io
+
+# Carga el modelo SavedModel
+model = tf.saved_model.load("efficientnet_alzheimer")
+
+# Define la función de preprocesamiento de imágenes
+def preprocess_image(image):
+  image = image.resize((200, 200)) 
+  image = np.array(image) / 255.0  
+  image = np.expand_dims(image, axis=0)  
+  return image
+
+def predict(image: Image.Image):
+    """
+    Realiza la predicción de la imagen.
+
+    Args:
+        image: Imagen en formato PIL.Image
+
+    Returns:
+        dict: Un diccionario con la predicción y la probabilidad.
+    """
+
+    image = preprocess_image(image)
+
+    # Realiza la predicción
+    predictions = model(image)
+
+    # Obtén la clase predicha y la probabilidad
+    predicted_class = np.argmax(predictions)
+    confidence = np.max(predictions)
+
+    # Mapea la clase numérica al nombre de la clase
+    class_names = {
+        0: "VeryMildDemented", 
+        1: "NonDemented", 
+        2: "ModerateDemented",
+        3: "MildDemented"
+    }
+    predicted_class_name = class_names.get(predicted_class, "Clase desconocida")
+
+    # Devuelve el resultado
+    return {"prediction": predicted_class_name, "confidence": confidence.item()}