Update app.py

app.py

@@ -30,10 +30,19 @@ def aggregate_predictions(predicted_probs, taxonomic_level, class_names):
         species = row['species']
         higher_level = row[taxonomic_level]
         
-        species_index = class_names.index(species)  # Index of the species in the prediction array
-        higher_level_index = unique_labels.index(higher_level)
-        
-        aggregated_predictions[:, higher_level_index] += predicted_probs[:, species_index]
+        if species in class_names:  # Check if species exists in class names
+            species_index = class_names.index(species)  # Index of the species in the prediction array
+
+            if higher_level in unique_labels:  # Check if higher level exists
+                higher_level_index = unique_labels.index(higher_level)
+                
+                # Only update if indices are valid
+                if species_index < predicted_probs.shape[1] and higher_level_index < aggregated_predictions.shape[1]:
+                    if predicted_probs[:, species_index].max() >= 0.80:  # Check confidence level
+                        aggregated_predictions[:, higher_level_index] += predicted_probs[:, species_index]
+                    else:
+                        # Stop aggregation at the current level if confidence is below 0.80
+                        break
     
     return aggregated_predictions, unique_labels