Spaces:
Sleeping
Sleeping
| 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 resnet import ResNet18 | |
| import gradio as gr | |
| model = ResNet18() | |
| model.load_state_dict(torch.load("model.pth", map_location=torch.device('cpu')), strict=False) | |
| 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] | |
| ) | |
| classes = ('plane', 'car', 'bird', 'cat', 'deer', | |
| 'dog', 'frog', 'horse', 'ship', 'truck') | |
| def resize_image_pil(image, new_width, new_height): | |
| # Convert to PIL image | |
| img = Image.fromarray(np.array(image)) | |
| # Get original size | |
| width, height = img.size | |
| # Calculate scale | |
| width_scale = new_width / width | |
| height_scale = new_height / height | |
| scale = min(width_scale, height_scale) | |
| # Resize | |
| resized = img.resize((int(width*scale), int(height*scale)), Image.NEAREST) | |
| # Crop to exact size | |
| resized = resized.crop((0, 0, new_width, new_height)) | |
| return resized | |
| def inference(input_img, transparency = 0.5, is_grad_cam=True, target_layer_number = -1, top_predictions=3): | |
| input_img = resize_image_pil(input_img, 32, 32) | |
| input_img = np.array(input_img) | |
| org_img = input_img | |
| input_img = input_img.reshape((32, 32, 3)) | |
| transform = transforms.ToTensor() | |
| input_img = transform(input_img) | |
| input_img = input_img | |
| input_img = input_img.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) | |
| if is_grad_cam: | |
| target_layers = [model.layer2[target_layer_number]] | |
| cam = GradCAM(model=model, target_layers=target_layers) | |
| grayscale_cam = cam(input_tensor=input_img, targets=None) | |
| grayscale_cam = grayscale_cam[0, :] | |
| visualization = show_cam_on_image(org_img/255, grayscale_cam, use_rgb=True, image_weight=transparency) | |
| else: | |
| visualization = None | |
| # Sort the confidences dictionary based on confidence values | |
| sorted_confidences = dict(sorted(confidences.items(), key=lambda item: item[1], reverse=True)) | |
| # Pick the top n predictions | |
| top_n_confidences = dict(list(sorted_confidences.items())[:top_predictions]) | |
| return classes[prediction[0].item()], visualization, top_n_confidences | |
| title = "CIFAR10 trained on ResNet18 Model with GradCAM" | |
| description = "A simple Gradio interface to infer on ResNet model, and get GradCAM results" | |
| examples = [["cat.jpg", 0.5, -1], ["dog.jpg", 0.5, -1]] | |
| demo = gr.Interface( | |
| inference, | |
| inputs = [ | |
| gr.Image(width=256, height=256, label="Input Image"), | |
| gr.Slider(0, 1, value = 0.5, label="Overall Opacity of Image"), | |
| gr.Checkbox(label="Show GradCAM"), | |
| gr.Slider(-2, -1, value = -2, step=1, label="Which Layer?"), | |
| gr.Slider(2, 10, value=3, step=1, label="Number of Top Classes") | |
| ], | |
| outputs = [ | |
| "text", | |
| gr.Image(width=256, height=256, label="Output"), | |
| gr.Label(label="Top Classes") | |
| ], | |
| title = title, | |
| description = description, | |
| examples = examples, | |
| ) | |
| demo.launch() |