Spaces:
Running
Running
import streamlit as st | |
import tensorflow as tf | |
import numpy as np | |
from keras.models import load_model | |
from tensorflow.keras.backend import clear_session | |
import cv2 | |
import os | |
st.set_page_config( | |
page_title = 'Patacotrón', | |
initial_sidebar_state = 'collapsed', | |
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' | |
} | |
) | |
col_a, col_b, = st.columns(2) | |
with col_a: | |
st.title("Entorno de ejecución") | |
st.caption("Los modelos no están en orden de eficacia, sino en orden de creación.") | |
# Get the absolute path to the current directory | |
current_dir = os.path.abspath(os.path.dirname(__file__)) | |
# Get the absolute path to the parent directory of the current directory | |
root_dir = os.path.abspath(os.path.join(current_dir, os.pardir)) | |
# Join the path to the models folder | |
DIR = os.path.join(root_dir, "models") | |
threshold = .8 | |
ultra_button = st.checkbox('Usar ensamblaje de los modelos con mayor eficacia hasta la fecha (mejores resultados)') | |
ultra_flag = False | |
if ultra_button: | |
ultra_flag = True | |
models = os.listdir(DIR) | |
model_dict = dict() | |
for model in models: | |
model_name = model.split(DIR) | |
model_name = str(model.split('.h5')[0]) | |
model_dir = os.path.join(DIR, model) | |
model_dict[model_name] = model_dir | |
ultraptctrn = ['ptctrn_v1.6', 'ptctrn_v1.8', 'ptctrn_v1.9.1', 'ptctrn_v1.12'] | |
# Create a dropdown menu to select the model | |
model_choice = st.multiselect("Seleccione uno o varios modelos de clasificación", model_dict.keys()) | |
selected_models = [] | |
def ensemble_model(model_list, img): | |
y_gorrito = np.zeros((1, 1)) | |
for model in model_list: | |
instance_model = load_model(model_dict[model]) | |
y_gorrito += float(instance_model.predict(np.expand_dims(img, 0))) | |
clear_session() | |
return y_gorrito/len(model_list) | |
for model in model_choice: | |
selected_models.append(model) | |
# Set the image dimensions | |
IMAGE_WIDTH = IMAGE_HEIGHT = 224 | |
uploaded_file = st.file_uploader(label = '',type= ['jpg','png', 'jpeg', 'jfif', 'webp', 'heic']) | |
with col_b: | |
if st.button('¿Hay un patacón en la imagen?'): | |
if uploaded_file is not None and len(selected_models) > 0 or ultra_flag: | |
# Load the image and resize it to the required dimensions | |
img = np.frombuffer(uploaded_file.read(), np.uint8) | |
img = cv2.imdecode(img, cv2.IMREAD_COLOR) | |
raw_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) | |
img = cv2.resize(img, (IMAGE_WIDTH, IMAGE_HEIGHT)) | |
# Convert the image to RGB and preprocess it for the model | |
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) | |
img = img / 255. | |
# Pass the image to the model and get the prediction | |
if ultra_flag: | |
with st.spinner('Cargando ultra-predicción...'): | |
y_gorrito = ensemble_model(ultraptctrn, img) | |
else: | |
with st.spinner('Cargando predicción...'): | |
y_gorrito = ensemble_model(selected_models, img) | |
if y_gorrito > threshold: | |
st.success("¡Patacón Detectado!") | |
else: | |
st.error("No se encontró rastro de patacón.") | |
st.caption(f'La probabilidad de que la imagen tenga un patacón es del: {round(float(y_gorrito), 2)*100}%') | |
st.image(raw_img) | |
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)') | |
else: | |
st.write('Revisa haber seleccionado los modelos y la imagen correctamente.') |