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

from transformers import TFAutoModel, AutoImageProcessor

backbone = TFAutoModel.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 = keras.models.load_model(
  "model.keras",
  custom_objects={
      "transpose_channels": transpose_channels,
      "deit_forward": deit_forward
  },
  safe_mode=False
)

# Clases
labels = ['glioma', 'meningioma', 'notumor', 'pituitary']

# Preprocesamiento (ajustado al modelo Deit-Tiny 224x224)
processor = AutoImageProcessor.from_pretrained("facebook/deit-tiny-patch16-224")

def predict(image):

    inputs = processor(images=image, return_tensors="np")
    pixel_values = inputs["pixel_values"][0]     # (3, 224, 224)

    pixel_values = np.transpose(pixel_values, (1, 2, 0))
    pixel_values = np.expand_dims(pixel_values, axis=0).astype(np.float32)

    preds = model.predict(pixel_values)[0]
    return {labels[i]: float(preds[i]) for i in range(len(labels))}

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()