EuroSat / app2.py
ismaeltorres00's picture
Update app2.py
3d6dcbd verified
raw
history blame contribute delete
954 Bytes
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()