Update utils/vulnerability_classifier.py

utils/vulnerability_classifier.py

@@ -28,7 +28,7 @@ label_dict= {0: 'Agricultural communities',
  16: 'Urban populations',
  17: 'Women and other genders'}
 
-def getlabels(preds):
+def get_vulnerability_labels(preds):
 
     """
     Function that takes the numerical predictions as an input and returns a list of the labels.
@@ -37,28 +37,11 @@ def getlabels(preds):
     
     # Get label names
     preds_list = preds.tolist()
+
+    # Get the name of the group where the prediction is equal to "1"
+    result = [label_dict[key] for key, value in enumerate(preds_list) if value == 1]
     
-    predictions_names=[]
-    
-    # loop through each prediction
-    for ele in preds_list:
-    
-      # see if there is a value 1 and retrieve index
-      try:
-        index_of_one = ele.index(1)
-      except ValueError:
-        index_of_one = "NA"
-    
-      # Retrieve the name of the label (if no prediction made = NA)
-      if index_of_one != "NA":
-        name  = label_dict[index_of_one]
-      else:
-        name = "Other"
-    
-      # Append name to list
-      predictions_names.append(name)
-    
-    return predictions_names
+    return result
         
 @st.cache_resource
 def load_vulnerabilityClassifier(config_file:str = None, classifier_name:str = None):