Update pages/Entorno de Ejecución.py

pages/Entorno de Ejecución.py

@@ -37,6 +37,11 @@ def preprocess(file_uploader, module = 'cv2'): #makes the uploaded image readabl
         img = Image.open(file_uploader)
     return img
 
+def multiclass_prediction(classifier): #made for hf zero-shot pipeline results
+    score = (max([classifier[i]['score'] for i in range(len(classifier))]))
+    labels = [predict['label'] for predict in classifier if score == predict['score']]
+    return (labels[0] if len(labels) == 1 else labels, score)
+
 @st.experimental_singleton
 def load_clip():
     classifier = pipeline("zero-shot-image-classification", model = 'openai/clip-vit-large-patch14-336')
@@ -206,11 +211,9 @@ with zero_shot:
                     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"]
+                    label, y_gorrito = multiclass_prediction(zs_classifier)
                         
-                    if round(float(y_gorrito*100)) >= threshold:
+                    if label == "A yellow deep fried smashed plantain":
                         st.success("¡Patacón Detectado!")
                     else:
                         st.error("No se considera que haya un patacón en la imagen")