Spaces:
Runtime error
Runtime error
| 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 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(255), | |
| transform.CenterCrop(224), | |
| transform.ToTensor(), | |
| ]) | |
| def transform_image(image, transforms): | |
| # img = cv2.cvtColor(cv2.imread(image_path), cv2.COLOR_BGR2RGB) | |
| img = transforms(image) | |
| img = img.unsqueeze(0) | |
| return img | |
| DenseNet = torchvision.models.densenet161(weights="DEFAULT") | |
| for param in DenseNet.parameters(): | |
| param.requires_grad = True | |
| in_features = DenseNet.classifier.in_features | |
| DenseNet.classifier = nn.Linear(in_features, 2) | |
| class ModelGradCam(nn.Module): | |
| def __init__(self, base_model): | |
| super(ModelGradCam, self).__init__() | |
| self.base_model = base_model | |
| self.features_conv = self.base_model.features | |
| self.pool = nn.AdaptiveAvgPool2d((1,1)) | |
| self.classifier = self.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, 2208) | |
| 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() | |