Update app.py

app.py

@@ -1,9 +1,9 @@
 # Import the libraries
 import numpy as np
 import pandas as pd
-from tensorflow.keras.models import load_model # type: ignore
-from tensorflow.keras.preprocessing.image import load_img, img_to_array # type: ignore
-from tensorflow.keras.applications.convnext import preprocess_input # type: ignore
+from tensorflow.keras.models import load_model
+from tensorflow.keras.preprocessing.image import load_img, img_to_array
+from tensorflow.keras.applications.convnext import preprocess_input
 import gradio as gr
 
 # Load the model
@@ -13,62 +13,83 @@ model = load_model('models/ConvNeXtBase_80_tresh_spp.tf')
 taxo_df = pd.read_csv('taxonomy/taxonomy_mapping.csv', sep=';')
 taxo_df['species'] = taxo_df['species'].str.replace('_', ' ')
 
-# Extract unique class names from the 'species' column
-class_names = sorted(taxo_df['species'].unique())
+# Available taxonomic levels
+taxonomic_levels = ['species', 'genus', 'family', 'order', 'class']
 
-# Function to map predicted class index to class name
-def get_class_name(predicted_class):
-    return class_names[predicted_class]
+# Function to map predicted class index to class name at the selected taxonomic level
+def get_class_name(predicted_class, taxonomic_level):
+    unique_labels = sorted(taxo_df[taxonomic_level].unique())
+    return unique_labels[predicted_class]
+
+# Function to aggregate predictions to the specified taxonomic level
+def aggregate_predictions(predicted_probs, taxonomic_level):
+    unique_labels = sorted(taxo_df[taxonomic_level].unique())
+    aggregated_predictions = np.zeros((predicted_probs.shape[0], len(unique_labels)))
+
+    for idx, row in taxo_df.iterrows():
+        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)  # Index of the higher taxonomic level
+        
+        aggregated_predictions[:, higher_level_index] += predicted_probs[:, species_index]
+    
+    return aggregated_predictions, unique_labels
 
 # Function to load and preprocess the image
 def load_and_preprocess_image(image, target_size=(224, 224)):
-    # Resize the image (assuming image is a PIL image)
+    # Resize the image
     img_array = img_to_array(image.resize(target_size))
-    # Expand the dimensions of the array to match model input
+    # Expand the dimensions to match model input
     img_array = np.expand_dims(img_array, axis=0)
-    # Preprocess using the appropriate function (for example, ResNet50)
+    # Preprocess the image
     img_array = preprocess_input(img_array)
     return img_array
 
-# Function to make predictions
-def make_prediction(image):
+# Function to make predictions at the selected taxonomic level
+def make_prediction(image, taxonomic_level):
     # Preprocess the image
     img_array = load_and_preprocess_image(image)
-    # Make a prediction
+    # Make a prediction at the species level
     prediction = model.predict(img_array)
     
+    # Aggregate predictions to the selected taxonomic level
+    aggregated_predictions, aggregated_class_labels = aggregate_predictions(prediction, taxonomic_level)
+    
     # Get the top 5 predictions
-    top_indices = np.argsort(prediction[0])[-5:][::-1]  # Get indices of top 5 classes
+    top_indices = np.argsort(aggregated_predictions[0])[-5:][::-1]  # Indices of top 5 predictions
+    
+    # Get the predicted class and common name for the top prediction
+    predicted_class_index = np.argmax(aggregated_predictions)
+    predicted_class_name = aggregated_class_labels[predicted_class_index]
     
-    # Get predicted class and common name for the top prediction
-    predicted_class_index = np.argmax(prediction)
-    predicted_class_name = get_class_name(predicted_class_index)
-    predicted_common_name = taxo_df[taxo_df['species'] == predicted_class_name]['common_name'].values[0]  # Get common name
-    confidence = prediction[0][predicted_class_index] * 100  # Confidence of the predicted class
+    # Get common name for the top predicted class
+    predicted_common_name = taxo_df[taxo_df[taxonomic_level] == predicted_class_name]['common_name'].values[0]
+    confidence = aggregated_predictions[0][predicted_class_index] * 100  # Confidence of the predicted class
     
     # Create output text with HTML formatting
-    output_text = f"<h1 style='font-weight: bold;'><span style='font-style: italic;'>{predicted_class_name}</span> ({predicted_common_name})</h1>"  # Large bold for predicted class, italic for class name
-    output_text += "<h4 style='font-weight: bold; font-size: 1.2em;'>Top 5 Predictions:</h4>"  # Bold and larger font for predictions
+    output_text = f"<h1 style='font-weight: bold;'><span style='font-style: italic;'>{predicted_class_name}</span> ({predicted_common_name})</h1>"
+    output_text += "<h4 style='font-weight: bold; font-size: 1.2em;'>Top 5 Predictions:</h4>"
     
     for i in top_indices:
-        class_name = get_class_name(i)
-        common_name = taxo_df[taxo_df['species'] == class_name]['common_name'].values[0]  # Get common name from CSV
-        confidence_percentage = prediction[0][i] * 100
-        
-        # Format the output with space between class name and common name
+        class_name = aggregated_class_labels[i]
+        common_name = taxo_df[taxo_df[taxonomic_level] == class_name]['common_name'].values[0]
+        confidence_percentage = aggregated_predictions[0][i] * 100
         output_text += f"<div style='display: flex; justify-content: space-between;'>" \
                        f"<span style='font-style: italic;'>{class_name}</span>&nbsp;(<span>{common_name}</span>)" \
                        f"<span style='margin-left: auto;'>{confidence_percentage:.2f}%</span></div>"
-
+    
     return output_text
 
 # Define the Gradio interface
 interface = gr.Interface(
-    fn=make_prediction,           # Function to be called for predictions
-    inputs=gr.Image(type="pil"),  # Input type: Image (PIL format)
-    outputs="html",               # Output type: HTML for formatting
-    title="Amazon arboreal species classification",
-    description="Upload an image to classify the species."
+    fn=make_prediction,            # Function to be called for predictions
+    inputs=[gr.Image(type="pil"),   # Input type: Image (PIL format)
+            gr.Dropdown(choices=taxonomic_levels, label="Taxonomic Level", default="species")],  # Dropdown for taxonomic level
+    outputs="html",                # Output type: HTML for formatting
+    title="Amazon Arboreal Species Classification",
+    description="Upload an image and select the taxonomic level to classify the species."
 )
 
 # Launch the Gradio interface