Spaces:
Sleeping
Sleeping
| import torch, torchvision | |
| from torchvision import transforms | |
| from torchvision import datasets | |
| import numpy as np | |
| import gradio as gr | |
| from PIL import Image | |
| from pytorch_grad_cam import GradCAM | |
| from pytorch_grad_cam.utils.image import preprocess_image, show_cam_on_image | |
| import itertools | |
| import matplotlib.pyplot as plt | |
| from utils import LitCIFAR10 | |
| model = LitCIFAR10.load_from_checkpoint("model.ckpt") | |
| model.eval() | |
| classes = ('plane', 'car', 'bird', 'cat', | |
| 'deer', 'dog', 'frog', 'horse', 'ship', 'truck') | |
| means = [0.4914, 0.4822, 0.4465] | |
| stds = [0.2470, 0.2435, 0.2616] | |
| cifar_testset = datasets.CIFAR10(root='.', train=False, download=True) | |
| transform=transforms.Compose([ | |
| transforms.ToTensor(), | |
| transforms.Normalize(means, stds) | |
| ]) | |
| class ClassifierOutputTarget: | |
| def __init__(self, category): | |
| self.category = category | |
| def __call__(self, model_output): | |
| if len(model_output.shape) == 1: | |
| return model_output[self.category] | |
| return model_output[:, self.category] | |
| def inference(wants_gradcam, n_gradcam, target_layer_number, transparency, wants_misclassified, n_misclassified, input_img = None, n_top_classes=10): | |
| if wants_gradcam: | |
| outputs_inference_gc = [] | |
| count_gradcam = 1 | |
| for data, target in cifar_testset: | |
| input_tensor = preprocess_image(data, | |
| mean=means, | |
| std=stds) | |
| target_layers = [model.model.layer3[target_layer_number]] | |
| targets = [ClassifierOutputTarget(target)] | |
| cam = GradCAM(model=model, target_layers=target_layers, use_cuda=False) | |
| grayscale_cam = cam(input_tensor=input_tensor, targets=targets) | |
| grayscale_cam = grayscale_cam[0, :] | |
| rgb_img = np.float32(data) / 255 | |
| visualization = np.array(show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True, image_weight=transparency)) | |
| outputs_inference_gc.append(visualization) | |
| count_gradcam += 1 | |
| if count_gradcam > n_gradcam: | |
| break | |
| else: | |
| outputs_inference_gc = None | |
| if wants_misclassified: | |
| outputs_inference_mis = [] | |
| count_mis = 1 | |
| for data_, target in cifar_testset: | |
| data = transform(data_) | |
| data = data.unsqueeze(0) | |
| output = model(data) | |
| pred = output.argmax(dim=1, keepdim=True) | |
| if pred.item()!=target: | |
| count_mis += 1 | |
| fig = plt.figure() | |
| fig.add_subplot(111) | |
| plt.imshow(data_) | |
| plt.title(f'Target: {classes[target]}\nPred: {classes[pred.item()]}') | |
| plt.axis('off') | |
| fig.canvas.draw() | |
| fig_img = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8) | |
| fig_img = fig_img.reshape(fig.canvas.get_width_height()[::-1] + (3,)) | |
| plt.close(fig) | |
| outputs_inference_mis.append(fig_img) | |
| if count_mis > n_misclassified: | |
| break | |
| else: | |
| outputs_inference_mis = None | |
| if input_img is not None: | |
| data = transform(input_img) | |
| data = data.unsqueeze(0) | |
| output = model(data) | |
| softmax = torch.nn.Softmax(dim=0) | |
| o = softmax(output.flatten()) | |
| confidences = {classes[i]: float(o[i]) for i in range(10)} | |
| _, prediction = torch.max(output, 1) | |
| confidences = {k: v for k, v in sorted(confidences.items(), key=lambda item: item[1], reverse=True)} | |
| confidences = dict(itertools.islice(confidences.items(), n_top_classes)) | |
| else: | |
| confidences = None | |
| return outputs_inference_gc, outputs_inference_mis, confidences | |
| title = "CIFAR10 trained on Custom ResNet Model with GradCAM" | |
| description = "A simple Gradio interface to infer on ResNet model, and get GradCAM results" | |
| examples = [[None, None, None, None, None, None, 'test_'+str(i)+'.jpg', None] for i in range(10)] | |
| demo = gr.Interface(inference, | |
| inputs = [gr.Checkbox(False, label='Do you want to see GradCAM outputs?'), | |
| gr.Slider(0, 10, value = 0, step=1, label="How many?"), | |
| gr.Slider(-2, -1, value = -2, step=1, label="Which target layer?"), | |
| gr.Slider(0, 1, value = 0, label="Opacity of GradCAM"), | |
| gr.Checkbox(False, label='Do you want to see misclassified images?'), | |
| gr.Slider(0, 10, value = 0, step=1, label="How many?"), | |
| gr.Image(shape=(32, 32), label="Input image"), | |
| gr.Slider(0, 10, value = 0, step=1, label="How many top classes you want to see?") | |
| ], | |
| outputs = [ | |
| gr.Gallery(label="GradCAM Outputs", show_label=True, elem_id="gallery").style(columns=[2], rows=[2], object_fit="contain", height="auto"), | |
| gr.Gallery(label="Misclassified Images", show_label=True, elem_id="gallery").style(columns=[2], rows=[2], object_fit="contain", height="auto"), | |
| gr.Label(num_top_classes=None) | |
| ], | |
| title = title, | |
| description = description, | |
| examples = examples | |
| ) | |
| demo.launch() |