Spaces:
Sleeping
Sleeping
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 |