import gradio as gr
from PIL import Image
from transformers import AutoModelForImageClassification, AutoFeatureExtractor
import torch

# Carga el modelo y el extractor de características de Hugging Face
model_id = "ismaeltorres00/ModeloFinalEuroSat"
model = AutoModelForImageClassification.from_pretrained(model_id)
feature_extractor = AutoFeatureExtractor.from_pretrained(model_id)

def predict_image(image):
    # Abre la imagen y convierte a RGB
    image = Image.open(image.name).convert("RGB")
    # Preprocesa la imagen
    inputs = feature_extractor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    # Obtén la predicción
    logits = outputs.logits
    predicted_class_idx = logits.argmax(-1).item()
    return model.config.id2label[predicted_class_idx]

iface = gr.Interface(fn=predict_image, inputs=gr.inputs.Image(type="file"), outputs="label", title="Image Prediction Model")
iface.launch()