wehub-resource-sync
86 lines
2.5 KiB
Python
86 lines
2.5 KiB
Python
import pytest
|
|
import torchvision
|
|
import torch
|
|
import cv2
|
|
from pytorch_grad_cam import GradCAM, \
|
|
ScoreCAM, \
|
|
GradCAMPlusPlus, \
|
|
AblationCAM, \
|
|
XGradCAM, \
|
|
EigenCAM, \
|
|
EigenGradCAM, \
|
|
LayerCAM, \
|
|
FullGrad, \
|
|
KPCA_CAM, \
|
|
SegEigenCAM
|
|
from pytorch_grad_cam.utils.image import show_cam_on_image, \
|
|
preprocess_image
|
|
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
|
|
|
|
|
|
@pytest.fixture
|
|
def numpy_image():
|
|
return cv2.imread("examples/both.png")
|
|
|
|
|
|
@pytest.mark.parametrize("cnn_model,target_layer_names", [
|
|
(torchvision.models.resnet18, ["layer4[-1]", "layer4[-2]"]),
|
|
(torchvision.models.vgg11, ["features[-1]"])
|
|
])
|
|
@pytest.mark.parametrize("batch_size,width,height", [
|
|
(2, 32, 32),
|
|
(1, 32, 40)
|
|
])
|
|
@pytest.mark.parametrize("target_category", [
|
|
None,
|
|
100
|
|
])
|
|
@pytest.mark.parametrize("aug_smooth", [
|
|
False
|
|
])
|
|
@pytest.mark.parametrize("eigen_smooth", [
|
|
True,
|
|
False
|
|
])
|
|
@pytest.mark.parametrize("cam_method",
|
|
[ScoreCAM,
|
|
AblationCAM,
|
|
GradCAM,
|
|
ScoreCAM,
|
|
GradCAMPlusPlus,
|
|
XGradCAM,
|
|
EigenCAM,
|
|
EigenGradCAM,
|
|
LayerCAM,
|
|
FullGrad,
|
|
KPCA_CAM,
|
|
SegEigenCAM])
|
|
|
|
def test_all_cam_models_can_run(numpy_image, batch_size, width, height,
|
|
cnn_model, target_layer_names, cam_method,
|
|
target_category, aug_smooth, eigen_smooth):
|
|
img = cv2.resize(numpy_image, (width, height))
|
|
input_tensor = preprocess_image(img)
|
|
input_tensor = input_tensor.repeat(batch_size, 1, 1, 1)
|
|
|
|
model = cnn_model(weights="DEFAULT")
|
|
target_layers = []
|
|
for layer in target_layer_names:
|
|
target_layers.append(eval(f"model.{layer}"))
|
|
|
|
cam = cam_method(model=model,
|
|
target_layers=target_layers)
|
|
cam.batch_size = 4
|
|
if target_category is None:
|
|
targets = None
|
|
else:
|
|
targets = [ClassifierOutputTarget(target_category)
|
|
for _ in range(batch_size)]
|
|
|
|
grayscale_cam = cam(input_tensor=input_tensor,
|
|
targets=targets,
|
|
aug_smooth=aug_smooth,
|
|
eigen_smooth=eigen_smooth)
|
|
assert(grayscale_cam.shape[0] == input_tensor.shape[0])
|
|
assert(grayscale_cam.shape[1:] == input_tensor.shape[2:])
|