Spaces:
Runtime error
Runtime error
File size: 4,115 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 118 119 120 121 122 123 |
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.transforms as transform
# !pip install efficientnet_pytorch -q
from efficientnet_pytorch import EfficientNet
if torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.device("cpu")
val_transform = transform.Compose([transform.Resize(size=(224, 224)),
transform.ToTensor(),
transform.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
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 Efficient(nn.Module):
def __init__(self, num_classes:int=1):
super(Efficient, self).__init__()
self.model = EfficientNet.from_pretrained("efficientnet-b3")
self.pool = nn.AdaptiveAvgPool2d((1,1))
self.fc = nn.Linear(1536, 256)
self.reg_model = nn.Sequential(
nn.BatchNorm1d(256),
nn.Linear(256, 500),
nn.BatchNorm1d(500),
nn.Tanh(),
nn.Dropout(0.2),
nn.Linear(500, 100),
nn.BatchNorm1d(100),
nn.Tanh(),
nn.Dropout(0.2),
nn.Linear(100, 4),
)
def forward(self, x):
x = self.model.extract_features(x)
x = self.pool(x)
x = x.view(-1, 1536)
x = self.fc(x)
x = self.reg_model(x)
return x
class ModelGradCam(nn.Module):
def __init__(self, base_model):
super(ModelGradCam, self).__init__()
self.base_model = base_model
self.features_conv = self.base_model.model.extract_features
self.pool = self.base_model.pool
self.fc = self.base_model.fc
self.classifier = self.base_model.reg_model
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, 1536)
x = self.fc(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, normalized=True):
model.eval()
# fig, axs = plt.subplots(1, 2, figsize=(15, 10))
image = x_ray_image
outputs = torch.nn.functional.softmax(model(image), dim=1)
_, pred = torch.max(outputs, 1)
outputs[0][pred.detach().cpu().numpy()[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.0025 + img
outputs = outputs.tolist()[0]
output_dict = dict(zip(class_names, np.round(outputs,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():.2f}')
# axs[0].axis('off')
# axs[1].set_title(f'Predicted: {class_names[pred.item()]}\n Confidence: {conf.item():.2f}')
# axs[1].axis('off')
# plt.show()
|