changes in ui

app.py

@@ -51,13 +51,22 @@ def load_model(model_name):
         extractor = AutoFeatureExtractor.from_pretrained(config["id"])
         model = AutoModelForImageClassification.from_pretrained(config["id"], output_attentions=True)
         model.eval()
+        # Enable gradients for class-specific attention
+        for param in model.parameters():
+            param.requires_grad = True
     elif config["type"] == "timm":
         model = timm.create_model(config["id"], pretrained=True)
         model.eval()
+        # Enable gradients for class-specific attention
+        for param in model.parameters():
+            param.requires_grad = True
         extractor = None
     elif config["type"] == "efficientnet":
         model = EfficientNet.from_pretrained(config["id"])
         model.eval()
+        # Enable gradients for class-specific attention
+        for param in model.parameters():
+            param.requires_grad = True
         extractor = None
 
     loaded_models[model_name] = (model, extractor)
@@ -129,6 +138,7 @@ def get_gradcam_for_class(model, image_tensor, class_idx):
 def vit_attention_for_class(model, extractor, image, class_idx):
     """Get attention map for specific class in ViT"""
     inputs = extractor(images=image, return_tensors="pt")
+    inputs['pixel_values'].requires_grad = True
     outputs = model(**inputs)
     
     score = outputs.logits[0, class_idx]
@@ -213,8 +223,39 @@ def vit_attention_rollout(outputs):
 
 
 # ---------------------------
-# Overlay Helper
+# Create Gradient Legend
 # ---------------------------
+def create_gradient_legend():
+    """Create a gradient legend image showing attention scale"""
+    width, height = 400, 60
+    gradient = np.zeros((height, width, 3), dtype=np.uint8)
+    
+    # Create gradient from blue to red (matching COLORMAP_JET)
+    for i in range(width):
+        # OpenCV's COLORMAP_JET: blue (low) -> cyan -> green -> yellow -> red (high)
+        value = int(255 * i / width)
+        color_single = np.array([[[value]]], dtype=np.uint8)
+        color_rgb = cv2.applyColorMap(color_single, cv2.COLORMAP_JET)
+        gradient[:, i] = color_rgb[0, 0]
+    
+    gradient = cv2.cvtColor(gradient, cv2.COLOR_BGR2RGB)
+    
+    # Convert to PIL and add text
+    from PIL import ImageDraw, ImageFont
+    gradient_pil = Image.fromarray(gradient)
+    draw = ImageDraw.Draw(gradient_pil)
+    
+    # Use default font
+    try:
+        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 14)
+    except:
+        font = ImageFont.load_default()
+    
+    # Add text labels
+    draw.text((10, 20), "Low Attention", fill=(255, 255, 255), font=font)
+    draw.text((width - 120, 20), "High Attention", fill=(255, 255, 255), font=font)
+    
+    return gradient_pil
 def overlay_attention(pil_img, attention_map):
     heatmap = (attention_map * 255).astype(np.uint8)
     heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
@@ -286,10 +327,10 @@ def predict(image, model_name, noise_level):
 def get_class_specific_attention(image, model_name, class_query):
     try:
         if image is None:
-            return None, "Please upload an image first"
+            return None, None, "Please upload an image first"
         
         if not class_query or class_query.strip() == "":
-            return None, "Please enter a class name"
+            return None, None, "Please enter a class name"
 
         # Find matching class
         class_query_lower = class_query.lower().strip()
@@ -307,7 +348,7 @@ def get_class_specific_attention(image, model_name, class_query):
                     break
             
             if matching_idx is None:
-                return None, f"Class '{class_query}' not found in model labels. Try a different class name or check suggestions."
+                return None, None, f"Class '{class_query}' not found in model labels. Try a different class name or check sample classes."
             
             # Get attention for this class
             att_map = vit_attention_for_class(model, extractor, image, matching_idx)
@@ -321,7 +362,7 @@ def get_class_specific_attention(image, model_name, class_query):
                     break
             
             if matching_idx is None:
-                return None, f"Class '{class_query}' not found in ImageNet labels. Try a different class name or check suggestions."
+                return None, None, f"Class '{class_query}' not found in ImageNet labels. Try a different class name or check sample classes."
             
             # Get Grad-CAM for this class
             transform = T.Compose([
@@ -335,13 +376,14 @@ def get_class_specific_attention(image, model_name, class_query):
             att_map = get_gradcam_for_class(model, x, matching_idx)
         
         overlay = overlay_attention(image, att_map)
-        return overlay, f"✓ Attention map generated for class: '{matched_label}'"
+        legend = create_gradient_legend()
+        return overlay, legend, f"✓ Attention map generated for class: '{matched_label}' (Index: {matching_idx})"
         
     except Exception as e:
         import traceback
         error_trace = traceback.format_exc()
         print(error_trace)
-        return None, f"Error generating attention map: {str(e)}"
+        return None, None, f"Error generating attention map: {str(e)}"
 
 
 # ---------------------------
@@ -400,16 +442,17 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 placeholder="e.g., cat, dog, car, pizza...",
                 info="Type any ImageNet class name"
             )
+            class_button = gr.Button("🎯 Generate Class-Specific Attention", variant="primary")
             gr.Markdown("**💡 Sample classes to try:**")
             sample_buttons = gr.Radio(
                 choices=SAMPLE_CLASSES,
                 label="Click to auto-fill",
                 interactive=True
             )
-            class_button = gr.Button("🎯 Generate Class-Specific Attention", variant="primary")
             
         with gr.Column(scale=2):
             class_output_image = gr.Image(label="🔍 Class-Specific Attention Map")
+            gradient_legend = gr.Image(label="📊 Attention Scale", show_label=True)
             class_status = gr.Textbox(label="Status", interactive=False)
     
     gr.Markdown("---")
@@ -438,7 +481,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     class_button.click(
         get_class_specific_attention,
         inputs=[input_image, model_dropdown, class_input],
-        outputs=[class_output_image, class_status]
+        outputs=[class_output_image, gradient_legend, class_status]
     )
 
 if __name__ == "__main__":