ALS-classification-validation/Data_augmentation_code.py at master · ex0pfe/ALS-classification-validation · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

!git clone https://github.com/ex0pfe/ALS-classifiaction-validation.git


!git clone https://github.com/YQ-XiaMLTech/ALS-classification.git


%cd ALS-classification

import sys
sys.path.append("model")

from model.DenseNet import SE_DenseNet
import torch
import PIL
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms, models
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
from model.SE_ResNet18 import SE_ResNet18
from process_dataset import pre_data
from torchvision.transforms.functional import to_pil_image
import os


model = SE_DenseNet(num_classes=2, dropout_rate=0.5)
!wget https://github.com/YQ-XiaMLTech/ALS-classification/raw/main/saves/model_fullDenseSE.pth -P /content/
model = torch.load('/content/model_fullDenseSE.pth', map_location='cpu', weights_only=False)

def apply_gradcam(input_tensor, model, target_layer, target_category):
    activations = None
    gradients = None

    def backward_hook(module, grad_input, grad_output):
        nonlocal gradients
        gradients = grad_output[0]

    def forward_hook(module, input, output):
        nonlocal activations
        activations = output

    hook_forward = target_layer.register_forward_hook(forward_hook)
    hook_backward = target_layer.register_full_backward_hook(backward_hook)

    output = model(input_tensor.unsqueeze(0))
    model.zero_grad()

    if target_category is None:
        target_category = output.argmax(dim=1)

    score = output[:, target_category].squeeze()
    score.backward()

    pooled_gradients = torch.mean(gradients, dim=[0, 2, 3])
    for i in range(activations.shape[1]):
        activations[:, i, :, :] *= pooled_gradients[i]

    heatmap = torch.mean(activations, dim=1).squeeze()
    heatmap = F.relu(heatmap)
    heatmap /= torch.max(heatmap)

    hook_forward.remove()
    hook_backward.remove()

    return heatmap

def main():

    img_folder_path = "../ALS-classifiaction-validation/data_augmentation/data_augmentation/color_denoising"  #Change data augmentation technique
    img_files = [f for f in os.listdir(img_folder_path) if f.endswith('.tif')]

    dataset_path = img_folder_path
    target_category = 2

    mean, std = pre_data.compute_mean_std(dataset_path)
    transform = transforms.Compose([
        transforms.Resize((400, 400)),
        transforms.ToTensor(),
        transforms.Normalize(mean=mean, std=std)
    ])

    for img_filename in img_files:
        img_path = os.path.join(img_folder_path, img_filename)

        img_original = Image.open(img_path)
        img = img_original.convert('RGB')
        img_tensor = transform(img)

        target_layer = model.features[-1]
        heatmap = apply_gradcam(img_tensor, model, target_layer, target_category)

        heatmap_np = heatmap.cpu().detach().numpy()
        heatmap_np = (heatmap_np - np.min(heatmap_np)) / (np.max(heatmap_np) - np.min(heatmap_np))
        heatmap_pil = to_pil_image(heatmap_np, mode='F').resize(img_original.size, PIL.Image.BICUBIC)

        overlay_np = np.array(heatmap_pil)

        mask_gradcam_70 = overlay_np > (0.7 * np.max(overlay_np))
        mask_gradcam_90 = overlay_np > (0.9 * np.max(overlay_np))

        overlay_colormap = cm.jet(overlay_np / np.max(overlay_np))
        overlay_colormap_rgb = (overlay_colormap[..., :3] * 255).astype(np.uint8)

        plt.figure(figsize=(15, 5))

        plt.subplot(1, 4, 1)
        plt.imshow(img_original)
        plt.title(f"Original: {img_filename}")
        plt.axis('off')

        plt.subplot(1, 4, 2)
        plt.imshow(img_original)
        plt.imshow(overlay_colormap_rgb, alpha=0.5)
        plt.title(f"Grad-CAM: {img_filename}")
        plt.axis('off')

        plt.subplot(1, 4, 3)
        plt.imshow(img_original)
        plt.imshow(mask_gradcam_70, cmap='gray', alpha=0.5)
        plt.title(f"Threshold 70%: {img_filename}")
        plt.axis('off')

        plt.subplot(1, 4, 4)
        plt.imshow(img_original)
        plt.imshow(mask_gradcam_90, cmap='gray', alpha=0.5)
        plt.title(f"Threshold 90%: {img_filename}")
        plt.axis('off')

        plt.show()

if __name__ == '__main__':
    main()