confidence value added

app.py

@@ -69,7 +69,6 @@ def load_model(model_name):
     loaded_models[model_name] = (model, extractor)
     return model, extractor
 
-
 # ---------------------------
 # Adversarial Noise
 # ---------------------------
@@ -79,7 +78,6 @@ def add_adversarial_noise(image, epsilon):
     noisy_img = np.clip(img_array + noise, 0, 1)
     return Image.fromarray((noisy_img * 255).astype(np.uint8))
 
-
 # ---------------------------
 # Grad-CAM for Class-Specific Attention
 # ---------------------------
@@ -126,7 +124,6 @@ def get_gradcam_for_class(model, image_tensor, class_idx):
 
     return cam
 
-
 # ---------------------------
 # ViT Attention for Class-Specific
 # ---------------------------
@@ -146,10 +143,10 @@ def vit_attention_for_class(model, extractor, image, class_idx):
         attn_map = attn.reshape(1, 14, 14)
         attn_map = attn_map.squeeze().detach().cpu().numpy()
         attn_map = (attn_map - attn_map.min()) / (attn_map.max() - attn_map.min() + 1e-8)
-        return attn_map
+        prob = F.softmax(outputs.logits, dim=-1)[0, class_idx].item()
+        return attn_map, prob
     
-    return np.ones((14, 14))
-
+    return np.ones((14, 14)), 0.0
 
 # ---------------------------
 # Grad-CAM Helper for CNNs
@@ -198,14 +195,12 @@ def get_gradcam(model, image_tensor):
 
     return cam
 
-
 # ---------------------------
 # ViT Attention Rollout
 # ---------------------------
 def vit_attention_rollout(outputs):
     if not hasattr(outputs, 'attentions') or outputs.attentions is None:
         return np.ones((14, 14))
-
     attn = outputs.attentions[-1]
     attn = attn.mean(1)
     attn = attn[:, 0, 1:]
@@ -214,7 +209,6 @@ def vit_attention_rollout(outputs):
     attn_map = (attn_map - attn_map.min()) / (attn_map.max() - attn_map.min() + 1e-8)
     return attn_map
 
-
 # ---------------------------
 # Overlay Attention on Image
 # ---------------------------
@@ -227,7 +221,6 @@ def overlay_attention(pil_img, attention_map):
     blended = Image.blend(pil_img.convert("RGB"), heatmap_pil, alpha=0.4)
     return blended
 
-
 # ---------------------------
 # Main Prediction Function
 # ---------------------------
@@ -280,9 +273,8 @@ def predict(image, model_name, noise_level):
         print(f"Error: {traceback.format_exc()}")
         return {"Error": str(e)}, None, None
 
-
 # ---------------------------
-# Class-Specific Attention
+# Class-Specific Attention with Confidence
 # ---------------------------
 def get_class_specific_attention(image, model_name, class_query):
     try:
@@ -295,6 +287,7 @@ def get_class_specific_attention(image, model_name, class_query):
         class_query_lower = class_query.lower().strip()
         matching_idx = None
         matched_label = None
+        confidence = 0.0
         
         model, extractor = load_model(model_name)
         
@@ -308,7 +301,7 @@ def get_class_specific_attention(image, model_name, class_query):
             if matching_idx is None:
                 return None, f"Class '{class_query}' not found in model labels."
             
-            att_map = vit_attention_for_class(model, extractor, image, matching_idx)
+            att_map, confidence = vit_attention_for_class(model, extractor, image, matching_idx)
             
         else:
             for idx, label in enumerate(IMAGENET_LABELS):
@@ -329,16 +322,18 @@ def get_class_specific_attention(image, model_name, class_query):
             x = transform(image).unsqueeze(0)
             x.requires_grad = True
             att_map = get_gradcam_for_class(model, x, matching_idx)
+            with torch.no_grad():
+                outputs = model(x)
+                confidence = F.softmax(outputs, dim=-1)[0, matching_idx].item()
         
         overlay = overlay_attention(image, att_map)
-        return overlay, f"✓ Attention map generated for class: '{matched_label}' (Index: {matching_idx})"
+        return overlay, f"✓ Attention map generated for class: '{matched_label}' (Index: {matching_idx}, Confidence: {confidence:.2f})"
         
     except Exception as e:
         import traceback
         print(traceback.format_exc())
         return None, f"Error generating attention map: {str(e)}"
 
-
 # ---------------------------
 # Sample Classes
 # ---------------------------
@@ -350,7 +345,6 @@ SAMPLE_CLASSES = [
     "person", "bicycle", "building", "tree", "flower"
 ]
 
-
 # ---------------------------
 # Gradio UI
 # ---------------------------