Update pages/Entorno de Ejecución.py

pages/Entorno de Ejecución.py

@@ -1,7 +1,7 @@
 import streamlit as st
 import tensorflow as tf
 from tensorflow.keras.models import load_model
-from transformers import pipeline
+from transformers import AutoConfig, AutoModel, pipeline
 from PIL import Image 
 import os
 import cv2
@@ -131,6 +131,14 @@ with vit:
     with col_a:
         st.title('Visual Transformers')
         st.caption('One class is all you need!')
+
+        model_dict = {
+            'google/vit-base-patch16-224-in21k' : 'frncscp/patacoptimus-prime',
+            'facebook/convnext-large-224' : 'frncscp/pataconxt'
+        }
+        
+        model_choice = st.multiselect("Seleccione un modelo de clasificación", model_dict.keys(), key = 'ViT_multiselect')
+        
         uploaded_file = st.file_uploader(key = 'ViT_upload', label = 'Sube la imagen a clasificar',type= ['jpg','png', 'jpeg', 'jfif', 'webp', 'heic'])
         flag = False
         threshold = st.slider('¿Cuál va a ser el límite desde donde se considere patacón? (se recomienda por encima del 80%)', 0, 100, 80, key = 'threshold_vit')
@@ -138,9 +146,11 @@ with vit:
     with col_b:
 
         if st.button(key = 'ViT_button', label ='¿Hay un patacón en la imagen?'):
-            if uploaded_file is not None:
+            if len(model_choice) != 1:
+                print('Recuerda seleccionar un solo modelo de clasificación')
+            elif uploaded_file is not None:
                 with st.spinner('Cargando predicción...'):
-                    classifier = pipeline("image-classification", model="frncscp/patacoptimus-prime")
+                    classifier = pipeline("image-classification", model= model_dict[model_choice[0]])
                     img = preprocess(uploaded_file, module = 'pil')
                 
                     classifier = classifier(img) 
@@ -165,24 +175,15 @@ with vit:
 with zero_shot:
     
     col_a, col_b = st.columns(2)
-
+        
     with col_a:
     
         st.title("Clasificación Zero-Shot")
-        st.caption("Usando Clip de OpenAI y Convnext de Facebook")
+        st.caption("Usando Clip de OpenAI")
 
         labels_for_classification =  ["A yellow deep fried smashed plantain", 
                             "Fried food", 
                             "Anything"]
-
-
-        model_dict = {
-            'Clip (openai/clip-vit-large-patch14-336)' : 'frncscp/patacoptimus-prime',
-            'Convnext (facebook/convnext-large-224)' : 'frncscp/pataconxt'
-        }
-        
-        model_choice = st.multiselect("Seleccione uno o varios modelos de clasificación", model_dict.keys(), key = 'zs_multiselect')
-
         
         uploaded_file = st.file_uploader(key = 'ZS_upload', label = 'Sube la imagen a clasificar',type= ['jpg','png', 'jpeg', 'jfif', 'webp', 'heic'])
         threshold = st.slider('¿Cuál va a ser el límite desde donde se considere patacón? (se recomienda por encima del 33.33%)', 0, 100, 34, key = 'threshold_ZS')
@@ -191,12 +192,10 @@ with zero_shot:
 
         
         if st.button(key = 'ZS_button', label ='¿Hay un patacón en la imagen?'):
-            if len(model_choice) != 1:
-                print('Recuerda seleccionar un solo modelo de clasificación')
-            elif uploaded_file is not None:
+            if uploaded_file is not None:
                 if not zsloaded:
                     with st.spinner("Cargando modelo de clasificación..."):
-                        classifier = pipeline("zero-shot-image-classification", model = model_dict[model_choice[0]])
+                        classifier = pipeline("zero-shot-image-classification", model = 'openai/clip-vit-large-patch14-336')
                         zsloaded = True
                 with st.spinner('Cargando predicción...'):
                     img = preprocess(uploaded_file, module = 'pil')