Update app.py

app.py

@@ -1,37 +1,17 @@
-import gradio as gr
-from transformers import ViTFeatureExtractor, ViTForImageClassification
-from PIL import Image
-import torch
+from transformers import ViTFeatureExtractor, AutoModelForImageClassification
 
 def classify_image(image):
-    # Cargar el feature extractor y el modelo
-    feature_extractor = ViTFeatureExtractor.from_pretrained("ismaeltorres00/ModeloFinalEuroSat")
-    model = ViTForImageClassification.from_pretrained("ismaeltorres00/ModeloFinalEuroSat")
+    try:
+        feature_extractor = ViTFeatureExtractor.from_pretrained("ismaeltorres00/ModeloFinalEuroSat")
+    except OSError as e:
+        # Manejo del error si el archivo no existe
+        print("No se pudo encontrar el archivo preprocessor_config.json. Verifica el repositorio.")
+        raise e
     
-    # Preprocesar la imagen
+    model = AutoModelForImageClassification.from_pretrained("ismaeltorres00/ModeloFinalEuroSat")
     inputs = feature_extractor(images=image, return_tensors="pt")
-    
-    # Mover el modelo a GPU si está disponible
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model.to(device)
-    inputs = {k: v.to(device) for k, v in inputs.items()}
-    
-    # Realizar la predicción
-    with torch.no_grad():
-        outputs = model(**inputs)
-    
-    # Obtener el label de la predicción
+    outputs = model(**inputs)
     logits = outputs.logits
-    predicted_class_idx = logits.argmax(-1).item()
-    return model.config.id2label[predicted_class_idx]
-
-# Crear la interfaz de Gradio
-demo = gr.Interface(
-    fn=classify_image, 
-    inputs=gr.Image(type="pil"), 
-    outputs="text",
-    title="Clasificación de Imágenes con ViT"
-)
 
-# Lanzar la aplicación
-demo.launch()
+    # Aquí procesarías los logits para obtener la clasificación
+    return logits.argmax(-1).item()