|
|
|
|
|
|
|
|
|
|
|
from io import BytesIO |
|
|
|
import matplotlib.pyplot as plt |
|
import requests |
|
import streamlit as st |
|
import torch |
|
from PIL import Image |
|
from torchvision import models |
|
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", |
|
"resnet50", |
|
"mobilenet_v3_small", |
|
"mobilenet_v3_large", |
|
"regnet_y_400mf", |
|
"convnext_tiny", |
|
"convnext_small", |
|
] |
|
LABEL_MAP = requests.get( |
|
"https://raw.githubusercontent.com/anishathalye/imagenet-simple-labels/master/imagenet-simple-labels.json" |
|
).json() |
|
|
|
|
|
def main(): |
|
|
|
|
|
st.set_page_config(layout="wide") |
|
|
|
|
|
st.title("TorchCAM: class activation explorer") |
|
|
|
st.write("\n") |
|
|
|
cols = st.columns((1, 1, 1)) |
|
cols[0].header("Input image") |
|
cols[1].header("Raw CAM") |
|
cols[-1].header("Overlayed CAM") |
|
|
|
|
|
|
|
st.sidebar.title("Input selection") |
|
|
|
st.set_option("deprecation.showfileUploaderEncoding", False) |
|
|
|
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") |
|
|
|
cols[0].image(img, use_column_width=True) |
|
|
|
|
|
st.sidebar.title("Setup") |
|
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..."): |
|
model = models.__dict__[tv_model](pretrained=True).eval() |
|
default_layer = locate_candidate_layer(model, (3, 224, 224)) |
|
|
|
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=[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) |
|
|
|
|
|
st.sidebar.write("\n") |
|
|
|
if st.sidebar.button("Compute CAM"): |
|
|
|
if uploaded_file is None: |
|
st.sidebar.error("Please upload an image first") |
|
|
|
else: |
|
with st.spinner("Analyzing..."): |
|
|
|
|
|
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() |
|
|
|
|
|
out = model(img_tensor.unsqueeze(0)) |
|
|
|
if class_selection == "Predicted class (argmax)": |
|
class_idx = out.squeeze(0).argmax().item() |
|
else: |
|
class_idx = LABEL_MAP.index(class_selection.rpartition(" - ")[-1]) |
|
|
|
act_maps = cam_extractor(class_idx, out) |
|
|
|
activation_map = act_maps[0] if len(act_maps) == 1 else cam_extractor.fuse_cams(act_maps) |
|
|
|
fig, ax = plt.subplots() |
|
ax.imshow(activation_map.squeeze(0).cpu().numpy()) |
|
ax.axis("off") |
|
cols[1].pyplot(fig) |
|
|
|
|
|
fig, ax = plt.subplots() |
|
result = overlay_mask(img, to_pil_image(activation_map, mode="F"), alpha=0.5) |
|
ax.imshow(result) |
|
ax.axis("off") |
|
cols[-1].pyplot(fig) |
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
