edit clas based atention

app.py

@@ -24,12 +24,10 @@ MODEL_CONFIGS = {
     "EfficientNet-B1": {"type": "efficientnet", "id": "efficientnet-b1"},
     "ResNet-50": {"type": "timm", "id": "resnet50"},
     "MobileNet-V2": {"type": "timm", "id": "mobilenetv2_100"},
-    "MobileNet-V3": {"type": "timm", "id": "mobilenetv3_small_100"},
     "MaxViT-Tiny": {"type": "timm", "id": "maxvit_tiny_tf_224"},
     "MobileViT-Small": {"type": "timm", "id": "mobilevit_s"},
     "EdgeNeXt-Small": {"type": "timm", "id": "edgenext_small"},
-    "RegNetY-002": {"type": "timm", "id": "regnety_002"},
-    "SqueezeNet": {"type": "timm", "id": "squeezenet1_1"}
+    "RegNetY-002": {"type": "timm", "id": "regnety_002"}
 }
 
 # ---------------------------
@@ -290,12 +288,13 @@ def get_class_specific_attention(image, model_name, class_query):
         if image is None:
             return None, "Please upload an image first"
         
-        if not class_query:
+        if not class_query or class_query.strip() == "":
             return None, "Please enter a class name"
 
         # Find matching class
         class_query_lower = class_query.lower().strip()
         matching_idx = None
+        matched_label = None
         
         model, extractor = load_model(model_name)
         
@@ -304,10 +303,11 @@ def get_class_specific_attention(image, model_name, class_query):
             for idx, label in model.config.id2label.items():
                 if class_query_lower in label.lower():
                     matching_idx = idx
+                    matched_label = label
                     break
             
             if matching_idx is None:
-                return None, f"Class '{class_query}' not found in model labels"
+                return None, f"Class '{class_query}' not found in model labels. Try a different class name or check suggestions."
             
             # Get attention for this class
             att_map = vit_attention_for_class(model, extractor, image, matching_idx)
@@ -317,10 +317,11 @@ def get_class_specific_attention(image, model_name, class_query):
             for idx, label in enumerate(IMAGENET_LABELS):
                 if class_query_lower in label.lower():
                     matching_idx = idx
+                    matched_label = label
                     break
             
             if matching_idx is None:
-                return None, f"Class '{class_query}' not found in ImageNet labels"
+                return None, f"Class '{class_query}' not found in ImageNet labels. Try a different class name or check suggestions."
             
             # Get Grad-CAM for this class
             transform = T.Compose([
@@ -334,37 +335,43 @@ 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 for class: {class_query}"
+        return overlay, f"✓ Attention map generated for class: '{matched_label}'"
         
     except Exception as e:
         import traceback
-        return None, f"Error: {str(e)}\n{traceback.format_exc()}"
+        error_trace = traceback.format_exc()
+        print(error_trace)
+        return None, f"Error generating attention map: {str(e)}"
 
 
 # ---------------------------
 # Get Class Suggestions
 # ---------------------------
 def get_class_suggestions(query, model_name):
-    if not query or len(query) < 2:
-        return []
-    
-    query_lower = query.lower()
-    suggestions = []
-    
-    model, extractor = load_model(model_name)
-    
-    if MODEL_CONFIGS[model_name]["type"] == "hf":
-        labels = list(model.config.id2label.values())
-    else:
-        labels = IMAGENET_LABELS
-    
-    for label in labels:
-        if query_lower in label.lower():
-            suggestions.append(label)
-            if len(suggestions) >= 10:  # Limit to 10 suggestions
-                break
-    
-    return suggestions
+    try:
+        if not query or len(query.strip()) < 2:
+            return gr.Dropdown(choices=[], value=None)
+        
+        query_lower = query.lower().strip()
+        suggestions = []
+        
+        model, extractor = load_model(model_name)
+        
+        if MODEL_CONFIGS[model_name]["type"] == "hf":
+            labels = list(model.config.id2label.values())
+        else:
+            labels = IMAGENET_LABELS
+        
+        for label in labels:
+            if query_lower in label.lower():
+                suggestions.append(label)
+                if len(suggestions) >= 10:  # Limit to 10 suggestions
+                    break
+        
+        return gr.Dropdown(choices=suggestions, value=None)
+    except Exception as e:
+        print(f"Error getting suggestions: {e}")
+        return gr.Dropdown(choices=[], value=None)
 
 
 # ---------------------------
@@ -372,7 +379,7 @@ def get_class_suggestions(query, model_name):
 # ---------------------------
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🧠 Enhanced Multi-Model Image Classifier")
-    gr.Markdown("### Features: Adversarial Examples | Class-Specific Attention | 15+ Models")
+    gr.Markdown("### Features: Adversarial Examples | Class-Specific Attention | 13+ Models")
     
     with gr.Row():
         with gr.Column(scale=1):
@@ -412,11 +419,12 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 info="Start typing to see suggestions"
             )
             class_suggestions = gr.Dropdown(
-                label="💡 Suggestions",
+                label="💡 Suggestions (click to use)",
                 choices=[],
-                interactive=True
+                interactive=True,
+                allow_custom_value=False
             )
-            class_button = gr.Button("🎯 Generate Class-Specific Attention")
+            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")
@@ -437,18 +445,21 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         outputs=[output_label, output_image, processed_image]
     )
     
+    # Update suggestions as user types
     class_input.change(
         get_class_suggestions,
         inputs=[class_input, model_dropdown],
         outputs=[class_suggestions]
     )
     
-    class_suggestions.change(
+    # When user selects from suggestions, update the text input
+    class_suggestions.select(
         lambda x: x,
         inputs=[class_suggestions],
         outputs=[class_input]
     )
     
+    # Generate attention map
     class_button.click(
         get_class_specific_attention,
         inputs=[input_image, model_dropdown, class_input],