Spaces:
Runtime error
Runtime error
| import torch, torchvision | |
| from torchvision import transforms | |
| 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 show_cam_on_image | |
| from torch.utils.data import DataLoader | |
| import itertools | |
| import matplotlib.pyplot as plt | |
| from custom_resnet import Custom_ResNet | |
| import utils as utils | |
| model = Custom_ResNet() | |
| model.load_state_dict(torch.load("results/custom_resnet_trained.pth", map_location=torch.device('cpu')), strict=False) | |
| model.eval() | |
| classes = ('plane', 'car', 'bird', 'cat', 'deer', | |
| 'dog', 'frog', 'horse', 'ship', 'truck') | |
| cifar_valid = utils.Cifar10SearchDataset('.', train=False, download=True, transform=utils.augmentation_custom_resnet('Valid')) | |
| 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] | |
| ) | |
| 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 = [] | |
| cifar_valid_loader = DataLoader(cifar_valid, batch_size=1, shuffle = True) | |
| count_gradcam = 1 | |
| for data, target in cifar_valid_loader: | |
| data, target = data.to('cpu'), target.to('cpu') | |
| target_layers = [model.layer2[target_layer_number]] | |
| cam = GradCAM(model=model, target_layers=target_layers, use_cuda=False) | |
| grayscale_cam = cam(input_tensor=data, targets=None) | |
| grayscale_cam = grayscale_cam[0, :] | |
| org_img = inv_normalize(data).squeeze(0).numpy() | |
| org_img = np.transpose(org_img, (1, 2, 0)) | |
| visualization = np.array(show_cam_on_image(org_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 = [] | |
| cifar_valid_loader = DataLoader(cifar_valid, batch_size=1, shuffle = True) | |
| count_mis = 1 | |
| for data, target in cifar_valid_loader: | |
| data, target = data.to('cpu'), target.to('cpu') | |
| outputs = model(data) | |
| softmax = torch.nn.Softmax(dim=0) | |
| o = softmax(outputs.flatten()) | |
| confidences = {classes[i]: float(o[i]) for i in range(10)} | |
| _, prediction = torch.max(outputs, 1) | |
| if target.numpy()[0] != prediction.numpy()[0]: | |
| count_mis += 1 | |
| org_img = inv_normalize(data).squeeze(0).numpy() | |
| org_img = np.transpose(org_img, (1, 2, 0)) | |
| fig = plt.figure() | |
| fig.add_subplot(111) | |
| plt.imshow(org_img) | |
| plt.title(f'Target: {classes[target.numpy()[0]]}\nPred: {classes[prediction.numpy()[0]]}') | |
| 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: | |
| transform=utils.augmentation_custom_resnet('Valid') | |
| org_img = input_img | |
| input_img = transform(image=input_img) | |
| input_img = input_img['image'].unsqueeze(0) | |
| outputs = model(input_img) | |
| softmax = torch.nn.Softmax(dim=0) | |
| o = softmax(outputs.flatten()) | |
| confidences = {classes[i]: float(o[i]) for i in range(10)} | |
| _, prediction = torch.max(outputs, 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, 'examples/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() |