Update pages/Entorno de Ejecución.py

pages/Entorno de Ejecución.py

@@ -37,10 +37,13 @@ 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
+def multiclass_prediction(classifier, important_class): #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)
+    for clase in classifier:
+        if clase['label'] == important_class:
+            class_score = clase['score']
+    return (labels[0] if len(labels) == 1 else labels, score, class_score)
 
 @st.cache_resource
 def load_clip():
@@ -227,7 +230,7 @@ with zero_shot:
                     zs_classifier = classifier(img, 
                     candidate_labels = labels_for_classification)
 
-                    label, y_gorrito = multiclass_prediction(zs_classifier)
+                    label, _, y_gorrito = multiclass_prediction(zs_classifier)
                         
                     if label == "A yellow deep fried smashed plantain":
                         st.success("¡Patacón Detectado!")