Spaces:
Runtime error
Runtime error
File size: 3,854 Bytes
49106b8 |
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 |
import cv2
from PIL import Image
import torch
import matplotlib.pyplot as plt
import torch.functional as F
import torch.nn as nn
import numpy as np
import torchvision
import torchvision.transforms as transforms
if torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.device("cpu")
mean_nums = [0.485, 0.456, 0.406]
std_nums = [0.229, 0.224, 0.225]
val_transform = transforms.Compose([
transforms.Resize((150,150)),
transforms.CenterCrop(150), #Performs Crop at Center and resizes it to 150x150
transforms.ToTensor(),
transforms.Normalize(mean=mean_nums, std = std_nums)
])
def transform_image(image, transforms):
# img = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB)
img = transforms(image)
img = img.unsqueeze(0)
return img
class DenseNet(nn.Module):
def __init__(self):
super(DenseNet, self).__init__()
self.base_model = torchvision.models.densenet121(weights="DEFAULT").features
self.pool = nn.AdaptiveAvgPool2d((1,1))
self.fc = nn.Linear(1024, 1000)
self.classify = nn.Linear(1000, 1)
self.classifier = nn.Sigmoid()
def forward(self, x):
x = self.base_model(x)
x = self.pool(x)
x = x.view(-1, 1024)
x = self.fc(x)
x = self.classify(x)
x = self.classifier(x)
return x
class ModelGradCam(nn.Module):
def __init__(self, base_model):
super(ModelGradCam, self).__init__()
self.features_conv = base_model.base_model
self.pool = base_model.pool
self.fc = base_model.fc
self.classify = base_model.classify
self.classifier = base_model.classifier
self.gradients = None
def activations_hook(self, grad):
self.gradients = grad
def forward(self, x):
x = self.features_conv(x)
h = x.register_hook(self.activations_hook)
x = self.pool(x)
x = x.view(-1, 1024)
x = self.fc(x)
x = self.classify(x)
x = self.classifier(x)
return x
def get_activations_gradient(self):
return self.gradients
def get_activations(self, x):
return self.features_conv(x)
def plot_grad_cam(model, x_ray_image, class_names, threshold:int=0.5, normalized=True):
model.eval()
# fig, axs = plt.subplots(1, 2, figsize=(15, 10))
image = x_ray_image
outputs = model(image).view(-1)
conf = [1-outputs.item(), outputs.item()]
# conf = 1 - outputs if outputs < threshold else outputs
pred = torch.where(outputs > threshold, torch.tensor(1, device=device), torch.tensor(0, device=device))
outputs[0].backward()
gradients = model.get_activations_gradient()
pooled_gradients = torch.mean(gradients, dim=[0, 2, 3])
activations = model.get_activations(image).detach()
activations *= pooled_gradients.unsqueeze(-1).unsqueeze(-1)
heatmap = torch.mean(activations, dim=1).squeeze()
heatmap = np.maximum(heatmap.cpu(), 0)
heatmap /= torch.max(heatmap)
img = image.squeeze().permute(1, 2, 0).cpu().numpy()
img = img if normalized else img/255.0
heatmap = cv2.resize(heatmap.numpy(), (img.shape[1], img.shape[0]))
heatmap = np.uint8(255 * heatmap)
heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
superimposed_img = heatmap * 0.0045 + img
output_dict = dict(zip(class_names, np.round(conf,3)))
return superimposed_img, class_names[pred.item()], output_dict
# axs[0].imshow(img)
# axs[1].imshow(superimposed_img)
# axs[0].set_title(f'Predicted: {class_names[pred.item()]}\n Confidence: {conf.item():.3f}')
# axs[0].axis('off')
# axs[1].set_title(f'Predicted: {class_names[pred.item()]}\n Confidence: {conf.item():.3f}')
# axs[1].axis('off')
# plt.show()
|