Spaces:
Sleeping
Sleeping
File size: 6,857 Bytes
9e78939 ee91973 165456f 9e78939 ee91973 9e78939 165456f ee91973 165456f ee91973 165456f ee91973 9e78939 5427919 9e78939 ee91973 ecb6f27 ee91973 9e78939 4eee11e 9e78939 165456f |
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 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 |
import torch.nn as nn
import torch.nn.functional as F
import torch
from torchvision import transforms
import cv2
import numpy as np
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image
import matplotlib.pyplot as plt
import textwrap
import io
def apply_normalization(chennels):
return nn.BatchNorm2d(chennels)
class CustomResnet(nn.Module):
def __init__(self):
super(CustomResnet, self).__init__()
# Input Block
drop = 0.0
# PrepLayer - Conv 3x3 s1, p1) >> BN >> RELU [64k]
self.preplayer = nn.Sequential(
nn.Conv2d(3, 64, (3, 3), padding=1, stride=1, bias=False), # 3
apply_normalization(64),
nn.ReLU(),
)
# Layer1 -
# X = Conv 3x3 (s1, p1) >> MaxPool2D >> BN >> RELU [128k]
self.convlayer1 = nn.Sequential(
nn.Conv2d(64, 128, (3, 3), padding=1, stride=1, bias=False), # 3
nn.MaxPool2d(2, 2),
apply_normalization(128),
nn.ReLU(),
)
# R1 = ResBlock( (Conv-BN-ReLU-Conv-BN-ReLU))(X) [128k]
self.reslayer1 = nn.Sequential(
nn.Conv2d(128, 128, (3, 3), padding=1, stride=1, bias=False), # 3
apply_normalization(128),
nn.ReLU(),
nn.Conv2d(128, 128, (3, 3), padding=1, stride=1, bias=False), # 3
apply_normalization(128),
nn.ReLU(),
)
# Conv 3x3 [256k]
self.convlayer2 = nn.Sequential(
nn.Conv2d(128, 256, (3, 3), padding=1, stride=1, bias=False), # 3
nn.MaxPool2d(2, 2),
apply_normalization(256),
nn.ReLU(),
)
# X = Conv 3x3 (s1, p1) >> MaxPool2D >> BN >> RELU [512k]
self.convlayer3 = nn.Sequential(
nn.Conv2d(256, 512, (3, 3), padding=1, stride=1, bias=False), # 3
nn.MaxPool2d(2, 2),
apply_normalization(512),
nn.ReLU(),
)
# R1 = ResBlock( (Conv-BN-ReLU-Conv-BN-ReLU))(X) [128k]
self.reslayer2 = nn.Sequential(
nn.Conv2d(512, 512, (3, 3), padding=1, stride=1, bias=False), # 3
apply_normalization(512),
nn.ReLU(),
nn.Conv2d(512, 512, (3, 3), padding=1, stride=1, bias=False), # 3
apply_normalization(512),
nn.ReLU(),
)
self.maxpool3 = nn.MaxPool2d(4, 2)
self.linear1 = nn.Linear(512,10)
def forward(self,x):
x = self.preplayer(x)
x1 = self.convlayer1(x)
x2 = self.reslayer1(x1)
x = x1+x2
x = self.convlayer2(x)
x = self.convlayer3(x)
x1 = self.reslayer2(x)
x = x+x1
x = self.maxpool3(x)
x = x.view(-1, 512)
x = self.linear1(x)
return F.log_softmax(x, dim=-1)
def resize_image(image, target_size=(200, 200)):
return cv2.resize(image, target_size)
def wrap_text(text, width=20):
return textwrap.fill(text, width)
import io
# define a function which returns an image as numpy array from figure
def get_img_from_fig(fig, dpi=180):
buf = io.BytesIO()
fig.savefig(buf, format="png", dpi=dpi)
buf.seek(0)
img_arr = np.frombuffer(buf.getvalue(), dtype=np.uint8)
buf.close()
img = cv2.imdecode(img_arr, 1)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
return img
def save_plot_as_image(images,texts, output_path):
num_images = len(images)
num_cols = min(4, num_images) # Assuming you want a maximum of 4 columns
num_rows = (num_images - 1) // num_cols + 1
fig, axes = plt.subplots(num_rows, num_cols, figsize=(3 * num_cols, 3 * num_rows))
subplot_height = 0.9 / num_rows # Adjust this value to control the height of each subplot
plt.subplots_adjust(hspace=subplot_height)
for i, ax in enumerate(axes.flat):
if i < num_images:
ax.imshow(images[i], cmap='gray')
ax.axis('off')
if texts is not None and i < len(texts):
wrapped_text = wrap_text(texts[i])
ax.set_title(wrapped_text, fontsize=12, pad=5)
else:
ax.axis('off')
plt.tight_layout()
# plt.savefig("tmp_arrays.png")
# plt.close()
return get_img_from_fig(plt)
# Function to run inference and return top classes
def get_gradcam(model,input_img, opacity,layer):
targets = None
inv_normalize = transforms.Normalize(
mean=[-0.50/0.23, -0.50/0.23, -0.50/0.23],
std=[1/0.23, 1/0.23, 1/0.23]
)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
transform = transforms.ToTensor()
input_img = transform(input_img)
input_img = input_img.to(device)
input_img = input_img.unsqueeze(0)
outputs = model(input_img)
if layer == "convblock1":
target_layers = model.convlayer1
elif layer == "convblock2":
target_layers = model.convlayer2
elif layer == "resblock1":
target_layers = model.reslayer1
elif layer == "resblock2":
target_layers = model.reslayer2
elif layer == "convblock3":
target_layers = model.convlayer3
layer_to_user = []
for i in target_layers:
if str(i) != "ReLU()":
layer_to_user.append(i)
print(layer_to_user)
final_outputs,texts = [],[]
for i in range(len(layer_to_user)):
cam = GradCAM(model=model, target_layers=[layer_to_user[i]], use_cuda=False)
grayscale_cam = cam(input_tensor=input_img, targets=targets)
grayscale_cam = grayscale_cam[0, :]
img = input_img.squeeze(0).to('cpu')
img = inv_normalize(img)
rgb_img = np.transpose(img, (1, 2, 0))
rgb_img = rgb_img.numpy()
visualization = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True, image_weight=opacity)
final_outputs.append(resize_image(visualization))
texts.append(str(layer_to_user[i]))
figure = save_plot_as_image(final_outputs,texts, "plot.png")
return figure
def get_misclassified_images(show_misclassified,num):
if show_misclassified:
return cv2.imread(f"missclassified_images_examples/{int(num)}.png")
else:
return None
def main_inference(num_of_output_classes,classes,model,input_img):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
transform = transforms.ToTensor()
input_img = transform(input_img)
input_img = input_img.to(device)
input_img = input_img.unsqueeze(0)
softmax = torch.nn.Softmax(dim=0)
outputs = model(input_img)
out = softmax(outputs.flatten())
_, prediction = torch.max(outputs, 1)
confidences = {classes[i]:float(out[i]) for i in range(num_of_output_classes)}
outputs = model(input_img)
_, prediction = torch.max(outputs, 1)
return confidences |