|
from pytorch_grad_cam import GradCAMPlusPlus |
|
from pytorch_grad_cam.utils.image import show_cam_on_image, preprocess_image |
|
import cv2 |
|
import numpy as np |
|
import torch |
|
import time |
|
import torch.nn as nn |
|
from configs import * |
|
import os, random |
|
|
|
|
|
model = MODEL.to(DEVICE) |
|
model.load_state_dict(torch.load(MODEL_SAVE_PATH, map_location=DEVICE)) |
|
model.eval() |
|
|
|
|
|
target_layer = None |
|
for child in model.features[-1]: |
|
if isinstance(child, nn.Conv2d): |
|
target_layer = child |
|
|
|
if target_layer is None: |
|
raise ValueError("Invalid layer name: {}".format(target_layer)) |
|
|
|
print(target_layer) |
|
|
|
def extract_gradcam(image_path=None, save_path=None): |
|
if image_path is None: |
|
for disease in CLASSES: |
|
print("Processing", disease) |
|
for image_path in os.listdir(r"data\test\Task 1\{}".format(disease)): |
|
print("Processing", image_path) |
|
image_path = r"data\test\Task 1\{}\{}".format(disease, image_path) |
|
image_name = image_path.split(".")[0].split("\\")[-1] |
|
rgb_img = cv2.imread(image_path, 1) |
|
rgb_img = np.float32(rgb_img) / 255 |
|
input_tensor = preprocess_image( |
|
rgb_img, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5] |
|
) |
|
input_tensor = input_tensor.to(DEVICE) |
|
input_tensor.requires_grad = True |
|
|
|
|
|
cam = GradCAMPlusPlus( |
|
model=model, target_layers=[target_layer], use_cuda=True |
|
) |
|
|
|
|
|
grayscale_cam = cam(input_tensor=input_tensor)[0] |
|
|
|
|
|
heatmap_colored = cv2.applyColorMap( |
|
np.uint8(255 * grayscale_cam), cv2.COLORMAP_JET |
|
) |
|
|
|
|
|
heatmap_colored = heatmap_colored.astype(np.float32) / 255 |
|
|
|
|
|
alpha = 0.3 |
|
|
|
|
|
final_output = cv2.addWeighted(rgb_img, 0.3, heatmap_colored, 0.7, 0) |
|
|
|
|
|
os.makedirs(f"docs/evaluation/gradcam/{disease}", exist_ok=True) |
|
cv2.imwrite( |
|
f"docs/evaluation/gradcam/{disease}/{image_name}.jpg", |
|
(final_output * 255).astype(np.uint8), |
|
) |
|
else: |
|
rgb_img = cv2.imread(image_path, 1) |
|
rgb_img = np.float32(rgb_img) / 255 |
|
input_tensor = preprocess_image( |
|
rgb_img, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5] |
|
) |
|
input_tensor = input_tensor.to(DEVICE) |
|
input_tensor.requires_grad = True |
|
|
|
|
|
cam = GradCAMPlusPlus(model=model, target_layers=[target_layer]) |
|
|
|
|
|
grayscale_cam = cam(input_tensor=input_tensor)[0] |
|
|
|
|
|
heatmap_colored = cv2.applyColorMap( |
|
np.uint8(255 * grayscale_cam), cv2.COLORMAP_JET |
|
) |
|
|
|
|
|
heatmap_colored = heatmap_colored.astype(np.float32) / 255 |
|
|
|
|
|
alpha = 0.3 |
|
|
|
|
|
final_output = cv2.addWeighted(rgb_img, 0.3, heatmap_colored, 0.7, 0) |
|
|
|
|
|
cv2.imwrite(save_path, (final_output * 255).astype(np.uint8)) |
|
|
|
return save_path |
|
|
|
|
|
|
|
|
|
|