Update pages/Entorno de Ejecución.py

pages/Entorno de Ejecución.py

@@ -2,6 +2,7 @@ import streamlit as st
 import tensorflow as tf
 from tensorflow.keras.models import load_model
 import os
+import cv2
 
 st.set_page_config(
     page_title = 'Patacotrón',
@@ -55,10 +56,19 @@ with cnn:
         @tf.function
         def predict(model_list, img):
             y_gorrito = 0
+            raw_img = tf.image.decode_image(img, channels=3)
+            img = tf.image.resize(raw_img,(IMAGE_WIDTH, IMAGE_HEIGHT))
             for model in model_list:
                 y_gorrito += tf.cast(model(tf.expand_dims(img/255., 0)), dtype=tf.float32)
-            return y_gorrito / len(model_list)
-        
+            return [y_gorrito / len(model_list), raw_img]
+
+        def preprocess(file_uploader): #converts it to .png
+            img = cv2.imdecode(bytearray(file_uploader.read()), cv2.IMREAD_COLOR)
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            img_tensor = tf.convert_to_tensor(img, dtype=tf.uint8)
+            png_bytes = tf.io.encode_png(img_tensor)
+            return png_bytes
+                
         # Set the image dimensions
         IMAGE_WIDTH = IMAGE_HEIGHT = 224
     
@@ -71,22 +81,33 @@ with cnn:
         if st.button('¿Hay un patacón en la imagen?'):
             if len(selected_models) > 0 and ultra_flag:
                 st.write('Debe elegir un solo método: Ultra-Patacotrón o selección múltiple.')
+            
             elif uploaded_file is not None:
-    
-                raw_img = tf.image.decode_image(uploaded_file.read(), channels=3)
-                img = tf.image.resize(raw_img,(IMAGE_WIDTH, IMAGE_HEIGHT))
-              
-                # Pass the image to the model and get the prediction
                 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
-                        y_gorrito = predict(ultraptctrn, img)
+                        try:
+                            y_gorrito, raw_img = predict(ultraptctrn, img)
+                        except:
+                            y_gorrito, raw_img = predict(ultraptctrn, preprocess(img))
+                        
+
+                # Pass the image to the model and get the prediction
+               # 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
+                        #y_gorrito = predict(ultraptctrn, 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]
-                        y_gorrito = predict(selected_models, img)
+                        try:             
+                            y_gorrito, raw_img = predict(selected_models, img)
+                        except:
+                            y_gorrito, raw_img = predict(selected_models, preprocess(img))                           
         
                 if round(float(y_gorrito*100)) >= threshold:
                     st.success("¡Patacón Detectado!")