Update app.py

app.py

@@ -6,7 +6,7 @@ import requests
 from io import BytesIO
 import numpy as np
 from pytorch_grad_cam import GradCAM, HiResCAM, ScoreCAM, GradCAMPlusPlus, AblationCAM, XGradCAM, EigenCAM, FullGrad
-from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget, get_target_layer
+from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
 from pytorch_grad_cam.utils.image import show_cam_on_image
 from timm.data import create_transform
 
@@ -51,7 +51,7 @@ def process_image(image_path, model):
     return tensor
 
 def get_cam_image(model, image, target_layer, cam_method):
-    cam = CAM_METHODS[cam_method](model=model, target_layers=[target_layer], use_cuda=torch.cuda.is_available())
+    cam = CAM_METHODS[cam_method](model=model, target_layers=[target_layer])
     grayscale_cam = cam(input_tensor=image)
     
     config = model.pretrained_cfg
@@ -69,14 +69,23 @@ def get_feature_info(model):
     else:
         return []
 
+def get_target_layer(model, target_layer_name):
+    if target_layer_name is None:
+        return None
+    
+    try:
+        return model.get_submodule(target_layer_name)
+    except AttributeError:
+        print(f"WARNING: Layer '{target_layer_name}' not found in the model.")
+        return None
+
 def explain_image(model_name, image_path, cam_method, feature_module):
     model = load_model(model_name)
     image = process_image(image_path, model)
     
-    if feature_module:
-        target_layer = get_target_layer(model, feature_module)
-        print(f"Using feature module: {feature_module}")
-    else:
+    target_layer = get_target_layer(model, feature_module)
+    
+    if target_layer is None:
         # Fallback to the last feature module or last convolutional layer
         feature_info = get_feature_info(model)
         if feature_info:
@@ -99,22 +108,29 @@ def explain_image(model_name, image_path, cam_method, feature_module):
 def update_feature_modules(model_name):
     model = load_model(model_name)
     feature_modules = get_feature_info(model)
-    return gr.Dropdown.update(choices=feature_modules, value=feature_modules[-1] if feature_modules else None)
+    return gr.Dropdown(choices=feature_modules, value=feature_modules[-1] if feature_modules else None)
 
-iface = gr.Interface(
-    fn=explain_image,
-    inputs=[
-        gr.Dropdown(choices=MODELS, label="Select Model"),
-        gr.Image(type="filepath", label="Upload Image"),
-        gr.Dropdown(choices=list(CAM_METHODS.keys()), label="Select CAM Method"),
-        gr.Dropdown(label="Select Feature Module (optional)")
-    ],
-    outputs=gr.Image(type="pil", label="Explained Image"),
-    title="Explainable AI with timm models",
-    description="Upload an image, select a model, CAM method, and optionally a specific feature module to visualize the explanation.",
-    allow_flagging="never"
-)
-
-iface.load(update_feature_modules, inputs=[iface.inputs[0]], outputs=[iface.inputs[3]])
+with gr.Blocks() as demo:
+    gr.Markdown("# Explainable AI with timm models")
+    gr.Markdown("Upload an image, select a model, CAM method, and optionally a specific feature module to visualize the explanation.")
+    
+    with gr.Row():
+        with gr.Column():
+            model_dropdown = gr.Dropdown(choices=MODELS, label="Select Model")
+            image_input = gr.Image(type="filepath", label="Upload Image")
+            cam_method_dropdown = gr.Dropdown(choices=list(CAM_METHODS.keys()), label="Select CAM Method")
+            feature_module_dropdown = gr.Dropdown(label="Select Feature Module (optional)")
+            explain_button = gr.Button("Explain Image")
+        
+        with gr.Column():
+            output_image = gr.Image(type="pil", label="Explained Image")
+    
+    model_dropdown.change(fn=update_feature_modules, inputs=[model_dropdown], outputs=[feature_module_dropdown])
+    
+    explain_button.click(
+        fn=explain_image,
+        inputs=[model_dropdown, image_input, cam_method_dropdown, feature_module_dropdown],
+        outputs=[output_image]
+    )
 
-iface.launch()
+demo.launch()