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 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) # Function to run inference and return top classes def get_gradcam(model,input_img, opacity): 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) _, prediction = torch.max(outputs, 1) target_layers = [model.convlayer3[-2]] cam = GradCAM(model=model, target_layers=target_layers, 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) return visualization def get_misclassified_images(show_misclassified,num): if show_misclassified: return cv2.imread(f"missclassified_images_examples/{int(num)}_missclassified.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 # def run_inference(input_img, num_of_output_classes,transparency): # 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)} # target_layers = [model.convlayer3[-2]] # cam = GradCAM(model=model, target_layers=target_layers, use_cuda=True) # 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=transparency) # return confidences, rgb_img, transparency,grayscale_cam