fix: Fixed matplotlib call

app.py

@@ -1,21 +1,22 @@
-# Copyright (C) 2020-2021, François-Guillaume Fernandez.
+# Copyright (C) 2021-2022, François-Guillaume Fernandez.
 
-# This program is licensed under the Apache License version 2.
-# See LICENSE or go to <https://www.apache.org/licenses/LICENSE-2.0.txt> for full license details.
+# This program is licensed under the Apache License 2.0.
+# See LICENSE or go to <https://www.apache.org/licenses/LICENSE-2.0> for full license details.
 
+from io import BytesIO
+
+import matplotlib.pyplot as plt
 import requests
 import streamlit as st
-import matplotlib.pyplot as plt
+import torch
 from PIL import Image
-from io import BytesIO
 from torchvision import models
-from torchvision.transforms.functional import resize, to_tensor, normalize, to_pil_image
+from torchvision.transforms.functional import normalize, resize, to_pil_image, to_tensor
 
 from torchcam import methods
 from torchcam.methods._utils import locate_candidate_layer
 from torchcam.utils import overlay_mask
 
-
 CAM_METHODS = ["CAM", "GradCAM", "GradCAMpp", "SmoothGradCAMpp", "ScoreCAM", "SSCAM", "ISCAM", "XGradCAM", "LayerCAM"]
 TV_MODELS = [
     "resnet18",
@@ -39,10 +40,7 @@ def main():
     # Designing the interface
     st.title("TorchCAM: class activation explorer")
     # For newline
-    st.write('\n')
-    st.write('Check the project at: https://github.com/frgfm/torch-cam')
-    # For newline
-    st.write('\n')
+    st.write("\n")
     # Set the columns
     cols = st.columns((1, 1, 1))
     cols[0].header("Input image")
@@ -53,36 +51,50 @@ def main():
     # File selection
     st.sidebar.title("Input selection")
     # Disabling warning
-    st.set_option('deprecation.showfileUploaderEncoding', False)
+    st.set_option("deprecation.showfileUploaderEncoding", False)
     # Choose your own image
-    uploaded_file = st.sidebar.file_uploader("Upload files", type=['png', 'jpeg', 'jpg'])
+    uploaded_file = st.sidebar.file_uploader("Upload files", type=["png", "jpeg", "jpg"])
     if uploaded_file is not None:
-        img = Image.open(BytesIO(uploaded_file.read()), mode='r').convert('RGB')
+        img = Image.open(BytesIO(uploaded_file.read()), mode="r").convert("RGB")
 
         cols[0].image(img, use_column_width=True)
 
     # Model selection
     st.sidebar.title("Setup")
-    tv_model = st.sidebar.selectbox("Classification model", TV_MODELS)
+    tv_model = st.sidebar.selectbox(
+        "Classification model",
+        TV_MODELS,
+        help="Supported models from Torchvision",
+    )
     default_layer = ""
     if tv_model is not None:
-        with st.spinner('Loading model...'):
+        with st.spinner("Loading model..."):
             model = models.__dict__[tv_model](pretrained=True).eval()
         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 torch.cuda.is_available():
+        model = model.cuda()
+
+    target_layer = st.sidebar.text_input(
+        "Target layer",
+        default_layer,
+        help='If you want to target several layers, add a "+" separator (e.g. "layer3+layer4")',
+    )
+    cam_method = st.sidebar.selectbox(
+        "CAM method",
+        CAM_METHODS,
+        help="The way your class activation map will be computed",
+    )
     if cam_method is not None:
         cam_extractor = methods.__dict__[cam_method](
-            model,
-            target_layer=target_layer.split("+") if len(target_layer) > 0 else None
+            model, target_layer=[s.strip() for s in 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)]
     class_selection = st.sidebar.selectbox("Class selection", ["Predicted class (argmax)"] + class_choices)
 
     # For newline
-    st.sidebar.write('\n')
+    st.sidebar.write("\n")
 
     if st.sidebar.button("Compute CAM"):
 
@@ -90,11 +102,14 @@ def main():
             st.sidebar.error("Please upload an image first")
 
         else:
-            with st.spinner('Analyzing...'):
+            with st.spinner("Analyzing..."):
 
                 # Preprocess image
                 img_tensor = normalize(to_tensor(resize(img, (224, 224))), [0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 
+                if torch.cuda.is_available():
+                    img_tensor = img_tensor.cuda()
+
                 # Forward the image to the model
                 out = model(img_tensor.unsqueeze(0))
                 # Select the target class
@@ -103,22 +118,22 @@ def main():
                 else:
                     class_idx = LABEL_MAP.index(class_selection.rpartition(" - ")[-1])
                 # Retrieve the CAM
-                cams = cam_extractor(class_idx, out)
+                act_maps = 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)
+                activation_map = act_maps[0] if len(act_maps) == 1 else cam_extractor.fuse_cams(act_maps)
                 # Plot the raw heatmap
                 fig, ax = plt.subplots()
-                ax.imshow(cam.numpy())
-                ax.axis('off')
+                ax.imshow(activation_map.squeeze(0).cpu().numpy())
+                ax.axis("off")
                 cols[1].pyplot(fig)
 
                 # Overlayed CAM
                 fig, ax = plt.subplots()
-                result = overlay_mask(img, to_pil_image(cam, mode='F'), alpha=0.5)
+                result = overlay_mask(img, to_pil_image(activation_map, mode="F"), alpha=0.5)
                 ax.imshow(result)
-                ax.axis('off')
+                ax.axis("off")
                 cols[-1].pyplot(fig)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()