pages/Entorno de Ejecución.py

import streamlit as st
import tensorflow as tf
from tensorflow.keras.models import load_model
from transformers import AutoConfig, AutoModel, pipeline
from PIL import Image 
import os
import cv2
import numpy as np

zsloaded = st.session_state.get('zsloaded', False)

st.set_page_config(
    page_title = 'Patacotrón',
    layout = 'wide',
    menu_items = {
        "About" : 'Proyecto ideado para la investigación de "Clasificación de imágenes de una sola clase con algortimos de Inteligencia Artificial".',
        "Report a Bug" : 'https://docs.google.com/forms/d/e/1FAIpQLScH0ZxAV8aSqs7TPYi86u0nkxvQG3iuHCStWNB-BoQnSW2V0g/viewform?usp=sf_link'
    }
)

st.sidebar.write("contact@patacotron.tech")

cnn, vit, zero_shot, autoencoder, svm, iforest, gan = st.tabs(["CNN", "ViT", "Zero-Shot", "Autoencoder", "OC-SVM", 'iForest', 'GAN'])

def predict(model_list, weights, img): #for non-supported formats
    y_gorrito = 0
    raw_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, (IMAGE_WIDTH, IMAGE_HEIGHT))
    for model, weight in zip(model_list, weights):
        y_gorrito += tf.cast(model(tf.expand_dims(img/255., 0)), dtype=tf.float32)*weight
    return [y_gorrito / sum(weights), raw_img]

def preprocess(file_uploader, module = 'cv2'): #makes the uploaded image readable
    img = np.frombuffer(uploaded_file.read(), np.uint8)
    if module == 'cv2':
        img = cv2.imdecode(img, cv2.IMREAD_COLOR)
    elif module == 'pil':
        img = Image.open(file_uploader)
    return img

with cnn:

    col_a, col_b, = st.columns(2)
    ultra_flag = None
    with col_a:
        st.title("Redes neuronales convolucionales")
        st.caption("Los modelos no están en orden de eficacia, sino en orden de creación.")
        
        current_dir = os.getcwd()
        root_dir = os.path.dirname(current_dir)
        
        # Join the path to the models folder
        DIR = os.path.join(current_dir, "models")
        models = os.listdir(DIR)
        common_root = r"/home/user/app/models/ptctrn_v"
        common_end = ".h5"
        
        model_dict = dict()
        for model in models: #preprocessing of strings so the name is readable in the multiselect bar
            model_dir = os.path.join(DIR, model)
            model_name = 'Patacotrón ' + model_dir.split(common_root)[-1].split(common_end)[0]
            model_dict[model_name] = model_dir
        #ultraversions = ['Patacotrón 1.5', 'Patacotrón 1.7', 'Patacotrón 1.8', 'Patacotrón 1.12', 'Patacotrón 1.12.2', 'Patacotrón 1.12.3']
        ultraversions = ['Patacotrón 1.5', 'Patacotrón 1.6', 'Patacotrón 1.12.2', 'Patacotrón 1.8', 'Patacotrón 1.12']#, 'Patacotrón 1.13.20', 'Patacotrón 1.13.38']
        #['Patacotrón 1.5', 'Patacotrón 1.6',  'Patacotrón 1.7', 'Patacotrón 1.12'] #
        ultra_button = st.checkbox('Ultra-Patacotrón (en construcción, no es la mejor versión)')
        ultra_flag = False
        weight_list = []
        
        if ultra_button:
            ultra_flag = True
            #weight_list = [3, 1, 4.5, 1, .8, 1] [.5, 1.75, 4, .5, 2]
            weight_list = [2.5, 1.8, 1.5, 3.14, 2.2] #.2, 2]
            #[1, 2, 3, 2.5] 
            st.caption('Para Ultra-Patacotrón, este porcentaje no representa una a priori una probabilidad, sino la combinación ponderada de modelos con sesgos positivos y negativos, lo importante es que identifique correctamente el objeto.')
            
        # Create a dropdown menu to select the model
        model_choice = st.multiselect("Seleccione uno o varios modelos de clasificación", model_dict.keys())
        
        threshold = st.slider('¿Cuál va a ser el límite donde se considere patacón? (el valor recomendado para Ultra-Patacotrón es 50%, para los demás, 75%-80%)', 0, 100, 50, key = 'threshold_convnet')
        
        selected_models = []
                
        # Set the image dimensions
        IMAGE_WIDTH = IMAGE_HEIGHT = 224
    
        executed = False
        
    with col_b:
             
        uploaded_file = st.file_uploader(key = 'convnet_upload', label = 'Sube la imagen a clasificar',type= ['jpg','png', 'jpeg', 'jfif', 'webp', 'heic'])
        
        if st.button(key = 'convnet_button', label ='¿Hay un patacón en la imagen?'):
            if (len(model_choice) > 0 and ultra_flag) or (len(model_choice) == 0 and ultra_flag is None):
                st.write('Debe elegir un método (debe seleccionar o deseleccionar alguno): Ultra-Patacotrón o Selección Múltiple.')
            
            elif uploaded_file is not None:
                img = preprocess(uploaded_file)
                if ultra_flag:
                    with st.spinner('Cargando ultra-predicción...'):
                        if not executed: 
                            ultraptctrn = [load_model(model_dict[model]) for model in ultraversions]
                            executed = True
                        final_weights = weight_list if len(weight_list) >= 1 else [1 for i in range(len(ultraptctrn))]
                        y_gorrito, raw_img = predict(ultraptctrn, final_weights, img)
                        
                else:
                    with st.spinner('Cargando predicción...'):
                        selected_models = [load_model(model_dict[model]) for model in model_choice if model not in selected_models]
                        final_weights = weight_list if len(weight_list) >= 1 else [1 for i in range(len(selected_models))]       
                        y_gorrito, raw_img = predict(selected_models, final_weights, img)

                if round(float(y_gorrito*100)) >= threshold:
                    st.success("¡Patacón Detectado!")
                
                else:
                    st.error("No se considera que haya un patacón en la imagen")
                
                st.caption(f'La probabilidad de que la imagen tenga un patacón es del: {round(float(y_gorrito * 100), 2)}%')  
                st.caption('Si los resultados no fueron los esperados, por favor, [haz click aquí](https://docs.google.com/forms/d/e/1FAIpQLScH0ZxAV8aSqs7TPYi86u0nkxvQG3iuHCStWNB-BoQnSW2V0g/viewform?usp=sf_link)')
                
                st.image(raw_img)
            
            else:
                st.write('Revisa haber seleccionado los modelos y la imagen correctamente.')

with vit:
      
    col_a, col_b = st.columns(2)

    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')

    with col_b:

        if st.button(key = 'ViT_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:
                with st.spinner('Cargando predicción...'):
                    classifier = pipeline("image-classification", model= model_dict[model_choice[0]])
                    img = preprocess(uploaded_file, module = 'pil')
                
                    classifier = classifier(img) 

                    for clase in classifier:
                        if clase['label'] == 'Patacon-True':
                            y_gorrito = clase["score"]
                
                #y_gorrito = classifier[0]["score"]
                    if round(float(y_gorrito*100)) >= threshold:
                        st.success("¡Patacón Detectado!")
                    else:
                        st.error("No se considera que haya un patacón en la imagen")
                    st.caption(f'La probabilidad de que la imagen tenga un patacón es del: {round(float(y_gorrito * 100), 2)}%')  
                    st.image(img) 
            else:
                st.write("Asegúrate de haber subido correctamente la imagen.")


                

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")

        labels_for_classification =  ["A yellow deep fried smashed plantain", 
                            "Fried food", 
                            "Anything"]
        
        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')
    
    with col_b:

        
        if st.button(key = 'ZS_button', label ='¿Hay un patacón en la imagen?'):
            if uploaded_file is not None:
                if not zsloaded:
                    with st.spinner("Cargando modelo de clasificación..."):
                        classifier = pipeline("zero-shot-image-classification", model = 'openai/clip-vit-large-patch14-336')
                        zsloaded = True
                        st.session_state['zsloaded'] = zsloaded
                with st.spinner('Cargando predicción...'):
                    img = preprocess(uploaded_file, module = 'pil')
                    zs_classifier = classifier(img, 
                    candidate_labels = labels_for_classification)

                    for clase in zs_classifier:
                        if clase['label'] == 'A yellow deep fried smashed plantain':
                            y_gorrito = clase["score"]
                        
                    if round(float(y_gorrito*100)) >= threshold:
                        st.success("¡Patacón Detectado!")
                    else:
                        st.error("No se considera que haya un patacón en la imagen")
                    st.caption(f'La probabilidad de que la imagen tenga un patacón es del: {round(float(y_gorrito * 100), 2)}%')  
                    st.image(img)
            else:
                st.write("Asegúrate de haber subido correctamente la imagen.")
        
with autoencoder:
    st.write('Próximamente')
with gan:
    st.write('Próximamente')  
with svm:
    st.write('Próximamente')
with iforest:
    st.write('Próximamente')