Upload app.py

app.py

@@ -0,0 +1,33 @@
+import streamlit as st
+import torch
+from transformers import ViTFeatureExtractor, ViTForImageClassification
+
+# Cargar el modelo preentrenado
+model_path = "output/model.bin"
+model = ViTForImageClassification.from_pretrained(model_path)
+
+# Cargar el extractor de características
+feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')
+
+# Función para hacer predicciones en una imagen de entrada
+def predict(image):
+    # Preprocesar la imagen
+    inputs = feature_extractor(image=image, return_tensors="pt")
+    # Hacer predicciones
+    outputs = model(**inputs)
+    # Obtener las etiquetas predichas
+    predicted_labels = torch.argmax(outputs.logits, dim=1)
+    # Devolver las etiquetas como una lista de strings
+    label_strings = ['inside', 'outside', 'food', 'drink', 'menu']
+    return [label_strings[label] for label in predicted_labels]
+
+# Interfaz de usuario para cargar una imagen y hacer predicciones
+st.title("ViT Image Classifier")
+uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
+if uploaded_file is not None:
+    image = Image.open(uploaded_file)
+    st.image(image, caption='Uploaded image.', use_column_width=True)
+    predictions = predict(image)
+    st.write("Predicted labels:")
+    for label in predictions:
+        st.write(label)