Update pages/Entorno de Ejecución.py

pages/Entorno de Ejecución.py

@@ -54,20 +54,25 @@ with cnn:
         selected_models = []
     
         @tf.function
-        def predict(model_list, img):
+        def tf_predict(model_list, img): #faster, but for few formats
             y_gorrito = 0
-            raw_img = tf.image.decode_image(img, channels=3)
-            img = tf.image.resize(raw_img,(IMAGE_WIDTH, IMAGE_HEIGHT))
+                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), 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
+        def basic_predict(model_list, img): #for non-supported formats
+            raw_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            img = cv2.resize(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), raw_img]
+
+        def preprocess(file_uploader): #makes the uploaded image readable
+            img = np.frombuffer(uploaded_file.read(), np.uint8)
+            img = cv2.imdecode(img, cv2.IMREAD_COLOR)
+            return img
                 
         # Set the image dimensions
         IMAGE_WIDTH = IMAGE_HEIGHT = 224
@@ -83,15 +88,16 @@ with cnn:
                 st.write('Debe elegir un solo método: 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
                         try:
-                            y_gorrito, raw_img = predict(ultraptctrn, uploaded_file)
+                            y_gorrito, raw_img = tf_predict(ultraptctrn, img)
                         except:
-                            y_gorrito, raw_img = predict(ultraptctrn, preprocess(uploaded_file))
+                            y_gorrito, raw_img = basic_predict(ultraptctrn, img)
                         
 
                 # Pass the image to the model and get the prediction
@@ -105,9 +111,9 @@ with cnn:
                     with st.spinner('Cargando predicción...'):
                         selected_models = [load_model(model_dict[model]) for model in model_choice if model not in selected_models]
                         try:             
-                            y_gorrito, raw_img = predict(selected_models, uploaded_file)
+                            y_gorrito, raw_img = tf_predict(selected_models, img)
                         except:
-                            y_gorrito, raw_img = predict(selected_models, preprocess(uploaded_file))                           
+                            y_gorrito, raw_img = basic_predict(selected_models, img)                           
         
                 if round(float(y_gorrito*100)) >= threshold:
                     st.success("¡Patacón Detectado!")