feat: Added option to retrieve CAMs from multiple layers

app.py

@@ -11,7 +11,8 @@ from io import BytesIO
 from torchvision import models
 from torchvision.transforms.functional import resize, to_tensor, normalize, to_pil_image
 
-from torchcam import cams
+from torchcam import methods
+from torchcam.methods._utils import locate_candidate_layer
 from torchcam.utils import overlay_mask
 
 
@@ -59,14 +60,14 @@ def main():
     if tv_model is not None:
         with st.spinner('Loading model...'):
             model = models.__dict__[tv_model](pretrained=True).eval()
-        default_layer = cams.utils.locate_candidate_layer(model, (3, 224, 224))
+        default_layer = locate_candidate_layer(model, (3, 224, 224))
 
     target_layer = st.sidebar.text_input("Target layer", default_layer)
     cam_method = st.sidebar.selectbox("CAM method", CAM_METHODS)
     if cam_method is not None:
-        cam_extractor = cams.__dict__[cam_method](
+        cam_extractor = methods.__dict__[cam_method](
             model,
-            target_layer=target_layer if len(target_layer) > 0 else None
+            target_layer=target_layer.split("+") if len(target_layer) > 0 else None
         )
 
     class_choices = [f"{idx + 1} - {class_name}" for idx, class_name in enumerate(LABEL_MAP)]
@@ -94,16 +95,18 @@ def main():
                 else:
                     class_idx = LABEL_MAP.index(class_selection.rpartition(" - ")[-1])
                 # Retrieve the CAM
-                activation_map = cam_extractor(class_idx, out)[0]
+                cams = cam_extractor(class_idx, out)
+                # Fuse the CAMs if there are several
+                cam = cams[0] if len(cams) == 1 else cam_extractor.fuse_cams(cams)
                 # Plot the raw heatmap
                 fig, ax = plt.subplots()
-                ax.imshow(activation_map.numpy())
+                ax.imshow(cam.numpy())
                 ax.axis('off')
                 cols[1].pyplot(fig)
 
                 # Overlayed CAM
                 fig, ax = plt.subplots()
-                result = overlay_mask(img, to_pil_image(activation_map, mode='F'), alpha=0.5)
+                result = overlay_mask(img, to_pil_image(cam, mode='F'), alpha=0.5)
                 ax.imshow(result)
                 ax.axis('off')
                 cols[-1].pyplot(fig)