Explanations: align families with classifier top-5; accept target_labels from UI

app.py

@@ -257,8 +257,34 @@ def generate_diagram_closest(input_image,model_name,top_k):
     diagram_path = get_diagram(embedding,top_k,model_name)
     return diagram_path
 
-def explain_image(input_image,model_name,explain_method,nb_samples,heatmap_alpha=0.22):
-    model,n_classes= get_model(model_name)
+def _top_k_classes_from_label_output(label_output, k=5):
+    """
+    Extract the top-k class names from a Gradio Label output.
+    Supports dict[label -> confidence] or list[(label, confidence), ...].
+    """
+    if not label_output:
+        return []
+    # Dict: sort by confidence, descending
+    if isinstance(label_output, dict):
+        items = sorted(label_output.items(), key=lambda kv: kv[1], reverse=True)
+        return [name for name, _ in items[:k]]
+    # List: assume list of (label, confidence) or labels
+    if isinstance(label_output, list):
+        names = []
+        for item in label_output[:k]:
+            if isinstance(item, (list, tuple)) and len(item) >= 1:
+                names.append(item[0])
+            else:
+                names.append(item)
+        return names
+    # Single label string
+    if isinstance(label_output, str):
+        return [label_output]
+    return []
+
+
+def explain_image(input_image, model_name, explain_method, nb_samples, class_predicted, heatmap_alpha=0.22):
+    model, n_classes = get_model(model_name)
     from explanations import explain
     if model_name in ('Fossils BEiT', 'Fossils 142'):
         size = 384
@@ -266,7 +292,19 @@ def explain_image(input_image,model_name,explain_method,nb_samples,heatmap_alpha
         size = 600
     #saliency, integrated, smoothgrad,
     h, w = input_image.shape[:2]
-    classes,exp_list = explain(model,input_image, h, w, explain_method,nb_samples,size = size, n_classes=n_classes, heatmap_alpha=heatmap_alpha)
+    target_labels = _top_k_classes_from_label_output(class_predicted, k=5)
+    classes, exp_list = explain(
+        model,
+        input_image,
+        h,
+        w,
+        explain_method,
+        nb_samples,
+        size=size,
+        n_classes=n_classes,
+        heatmap_alpha=heatmap_alpha,
+        target_labels=target_labels or None,
+    )
     #original =  saliency + integrated + smoothgrad 
     print('done')
     
@@ -959,8 +997,8 @@ with gr.Blocks(theme='sudeepshouche/minimalist', css=custom_css) as demo:
         #         classify_button = gr.Button("Classify Image")
 
 
-        def update_exp_outputs(input_image,model_name,explain_method,nb_samples,heatmap_alpha):
-            labels, images = explain_image(input_image,model_name,explain_method,nb_samples,heatmap_alpha)
+        def update_exp_outputs(input_image, model_name, explain_method, nb_samples, class_predicted, heatmap_alpha):
+            labels, images = explain_image(input_image, model_name, explain_method, nb_samples, class_predicted, heatmap_alpha)
             #labels_html = "".join([f'<div style="display: inline-block; text-align: center; width: 18%;">{label}</div>' for label in labels])
             #labels_markdown = f"<div style='width: 100%; text-align: center;'>{labels_html}</div>"
             image_caption=[]
@@ -968,7 +1006,11 @@ with gr.Blocks(theme='sudeepshouche/minimalist', css=custom_css) as demo:
                 image_caption.append((images[i],"Predicted Plant Family "+str(i)+": "+labels[i]))
             return image_caption
 
-        generate_explanations.click(fn=update_exp_outputs, inputs=[original_image,model_name,explain_method,sampling_size,heatmap_alpha], outputs=[exp_gallery])
+        generate_explanations.click(
+            fn=update_exp_outputs,
+            inputs=[original_image, model_name, explain_method, sampling_size, class_predicted, heatmap_alpha],
+            outputs=[exp_gallery],
+        )
 
         #find_closest_btn.click(find_closest, inputs=[input_image,model_name], outputs=[label_closest_image_0,label_closest_image_1,label_closest_image_2,label_closest_image_3,label_closest_image_4,closest_image_0,closest_image_1,closest_image_2,closest_image_3,closest_image_4])
         def _closest_table_html(labels, images, filenames):

explanations.py

@@ -1,8 +1,22 @@
 import xplique
 import tensorflow as tf
-from xplique.attributions import (Saliency, GradientInput, IntegratedGradients, SmoothGrad, VarGrad,
-                                  SquareGrad, GradCAM, Occlusion, Rise, GuidedBackprop,
-                                  GradCAMPP, Lime, KernelShap,SobolAttributionMethod,HsicAttributionMethod)
+from xplique.attributions import (
+    Saliency,
+    GradientInput,
+    IntegratedGradients,
+    SmoothGrad,
+    VarGrad,
+    SquareGrad,
+    GradCAM,
+    Occlusion,
+    Rise,
+    GuidedBackprop,
+    GradCAMPP,
+    Lime,
+    KernelShap,
+    SobolAttributionMethod,
+    HsicAttributionMethod,
+)
 from xplique.attributions.global_sensitivity_analysis import LatinHypercube
 import numpy as np
 import matplotlib.pyplot as plt
@@ -10,6 +24,7 @@ from inference_resnet import inference_resnet_finer, preprocess
 from labels import lookup_140
 import cv2
 BATCH_SIZE = 1
+_FAMILY_TO_INDEX = {v: k for k, v in lookup_140.items()}
 
 
 def letterbox_preprocess(img, size):
@@ -94,7 +109,18 @@ def show(img, original_size, output_size,p=False, **kwargs):
     
 
 
-def explain(model, input_image,h,w,explain_method,nb_samples,size=600, n_classes=171, heatmap_alpha=0.22) :
+def explain(
+    model,
+    input_image,
+    h,
+    w,
+    explain_method,
+    nb_samples,
+    size=600,
+    n_classes=171,
+    heatmap_alpha=0.22,
+    target_labels=None,
+):
     """
     Generate explanations for a given model and dataset.
     :param model: The model to explain.
@@ -142,12 +168,27 @@ def explain(model, input_image,h,w,explain_method,nb_samples,size=600, n_classes
     content_mask = np.zeros((size, size), dtype=np.float32)
     content_mask[content_top : content_top + content_h, content_left : content_left + content_w] = 1.0
 
-    predictions = class_model.predict(np.array([X]))
-    #Y = np.argmax(predictions)
-    top_5_indices = np.argsort(predictions[0])[-5:][::-1]
+    # Determine which classes to explain:
+    # - If target_labels are provided (from classifier output), use those.
+    # - Otherwise, fall back to top-5 classes from this forward pass.
     classes = []
-    for index in top_5_indices:
-        classes.append(lookup_140[index])
+    if target_labels:
+        indices = []
+        for name in target_labels:
+            idx = _FAMILY_TO_INDEX.get(name)
+            if idx is not None:
+                indices.append(idx)
+        if indices:
+            top_5_indices = np.array(indices, dtype=int)
+            classes = [lookup_140[i] for i in top_5_indices]
+        else:
+            predictions = class_model.predict(np.array([X]))
+            top_5_indices = np.argsort(predictions[0])[-5:][::-1]
+            classes = [lookup_140[i] for i in top_5_indices]
+    else:
+        predictions = class_model.predict(np.array([X]))
+        top_5_indices = np.argsort(predictions[0])[-5:][::-1]
+        classes = [lookup_140[i] for i in top_5_indices]
     #print(top_5_indices)
     X = np.expand_dims(X, 0)
     explanations = []