import gradio as gr
import tensorflow as tf
from PIL import Image
import numpy as np

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, :]

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)
IMG_SIZE = (224, 224)

def predict(image):
    # Convertimos a array
    img = image.resize(IMG_SIZE)
    img = np.array(img) / 255.0
    img = np.expand_dims(img, axis=0)

    # Predicción
    preds = model.predict(img)[0]
    probs = {labels[i]: float(preds[i]) for i in range(len(labels))}

    return probs

interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Sube una imagen"),
    outputs=gr.Label(num_top_classes=4, label="Predicción"),
    title="Clasificador de Tumores Cerebrales",
    description="Clasificador entrenado en 4 categorías: pituitary, meningioma, notumor, glioma"
)

interface.launch()