Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| from models.custom_resnet import CustomResNet | |
| from modules.visualize import plot_gradcam_images, plot_misclassified_images | |
| from pytorch_grad_cam import GradCAM | |
| from pytorch_grad_cam.utils.image import show_cam_on_image | |
| from torchvision import transforms | |
| import modules.config as config | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| TITLE = "CIFAR10 Image classification using a Custom ResNet Model" | |
| DESCRIPTION = "Gradio App to infer using a Custom ResNet model and get GradCAM results" | |
| examples = [ | |
| ["assets/images/airplane.jpg", 3, True, "layer3_x", 0.6, True, 5, True, 5], | |
| ["assets/images/bird.jpeg", 4, True, "layer3_x", 0.7, True, 10, True, 20], | |
| ["assets/images/car.jpg", 5, True, "layer3_x", 0.5, True, 15, True, 5], | |
| ["assets/images/cat.jpeg", 6, True, "layer3_x", 0.65, True, 20, True, 10], | |
| ["assets/images/deer.jpg", 7, False, "layer2", 0.75, True, 5, True, 5], | |
| ["assets/images/dog.jpg", 8, True, "layer2", 0.55, True, 10, True, 5], | |
| ["assets/images/frog.jpeg", 9, True, "layer2", 0.8, True, 15, True, 15], | |
| ["assets/images/horse.jpg", 10, False, "layer1_r1", 0.85, True, 20, True, 5], | |
| ["assets/images/ship.jpg", 3, True, "layer1_r1", 0.4, True, 5, True, 15], | |
| ["assets/images/truck.jpg", 4, True, "layer1_r1", 0.3, True, 5, True, 10], | |
| ] | |
| # load and initialise the model | |
| model = CustomResNet() | |
| # Define the device | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| # Using the checkpoint path present in config, load the trained model | |
| model.load_state_dict(torch.load(config.GRADIO_MODEL_PATH, map_location=device), strict=False) | |
| # Send model to CPU | |
| model.to(device) | |
| # Make the model in evaluation mode | |
| model.eval() | |
| # Load the misclassified images data | |
| misclassified_image_data = torch.load(config.GRADIO_MISCLASSIFIED_PATH, map_location=device) | |
| # Class Names | |
| classes = list(config.CIFAR_CLASSES) | |
| # Allowed model names | |
| model_layer_names = ["prep", "layer1_x", "layer1_r1", "layer2", "layer3_x", "layer3_r2"] | |
| def get_target_layer(layer_name): | |
| """Get target layer for visualization""" | |
| if layer_name == "prep": | |
| return [model.prep[-1]] | |
| elif layer_name == "layer1_x": | |
| return [model.layer1_x[-1]] | |
| elif layer_name == "layer1_r1": | |
| return [model.layer1_r1[-1]] | |
| elif layer_name == "layer2": | |
| return [model.layer2[-1]] | |
| elif layer_name == "layer3_x": | |
| return [model.layer3_x[-1]] | |
| elif layer_name == "layer3_r2": | |
| return [model.layer3_r2[-1]] | |
| else: | |
| return None | |
| def generate_prediction(input_image, num_classes=3, show_gradcam=True, transparency=0.6, layer_name="layer3_x"): | |
| """ "Given an input image, generate the prediction, confidence and display_image""" | |
| mean = list(config.CIFAR_MEAN) | |
| std = list(config.CIFAR_STD) | |
| transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean, std)]) | |
| with torch.no_grad(): | |
| orginal_img = input_image | |
| input_image = transform(input_image).unsqueeze(0).to(device) | |
| # print(f"Input Device: {input_image.device}") | |
| model_output = model(input_image).to(device) | |
| # print(f"Output Device: {outputs.device}") | |
| output_exp = torch.exp(model_output).to(device) | |
| # print(f"Output Exp Device: {o.device}") | |
| output_numpy = np.squeeze(np.asarray(output_exp.numpy())) | |
| # get indexes of probabilties in descending order | |
| sorted_indexes = np.argsort(output_numpy)[::-1] | |
| # sort the probabilities in descending order | |
| # final_class = classes[o_np.argmax()] | |
| confidences = {} | |
| for _ in range(int(num_classes)): | |
| # set the confidence of highest class with highest probability | |
| confidences[classes[sorted_indexes[_]]] = float(output_numpy[sorted_indexes[_]]) | |
| # Show Grad Cam | |
| if show_gradcam: | |
| # Get the target layer | |
| target_layers = get_target_layer(layer_name) | |
| cam = GradCAM(model=model, target_layers=target_layers) | |
| cam_generated = cam(input_tensor=input_image, targets=None) | |
| cam_generated = cam_generated[0, :] | |
| display_image = show_cam_on_image(orginal_img / 255, cam_generated, use_rgb=True, image_weight=transparency) | |
| else: | |
| display_image = orginal_img | |
| return confidences, display_image | |
| def app_interface( | |
| input_image, | |
| num_classes, | |
| show_gradcam, | |
| layer_name, | |
| transparency, | |
| show_misclassified, | |
| num_misclassified, | |
| show_gradcam_misclassified, | |
| num_gradcam_misclassified, | |
| ): | |
| """Function which provides the Gradio interface""" | |
| input_image = resize_image_pil(input_image, 32, 32) | |
| input_image = np.array(input_image) | |
| org_img = input_image | |
| # Get the prediction for the input image along with confidence and display_image | |
| confidences, display_image = generate_prediction(org_img, num_classes, show_gradcam, transparency, layer_name) | |
| if show_misclassified: | |
| misclassified_fig, misclassified_axs = plot_misclassified_images( | |
| data=misclassified_image_data, class_label=classes, num_images=num_misclassified | |
| ) | |
| else: | |
| misclassified_fig = None | |
| if show_gradcam_misclassified: | |
| gradcam_fig, gradcam_axs = plot_gradcam_images( | |
| model=model, | |
| data=misclassified_image_data, | |
| class_label=classes, | |
| # Use penultimate block of resnet18 layer 3 as the target layer for gradcam | |
| # Decided using model summary so that dimensions > 7x7 | |
| target_layers=get_target_layer(layer_name), | |
| targets=None, | |
| num_images=num_gradcam_misclassified, | |
| image_weight=transparency, | |
| ) | |
| else: | |
| gradcam_fig = None | |
| # # delete ununsed axises | |
| # del misclassified_axs | |
| # del gradcam_axs | |
| return confidences, display_image, misclassified_fig, gradcam_fig | |
| 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 | |
| inference_app = gr.Interface( | |
| app_interface, | |
| inputs=[ | |
| # This accepts the image after resizing it to 32x32 which is what our model expects | |
| gr.Image(width=256, height=256, label="Input Image"), | |
| gr.Number(value=3, maximum=10, minimum=1, step=1.0, precision=0, label="#Classes to show"), | |
| gr.Checkbox(True, label="Show GradCAM Image"), | |
| gr.Dropdown(model_layer_names, value="layer3_x", label="Visulalization Layer from Model"), | |
| # How much should the image be overlayed on the original image | |
| gr.Slider(0, 1, 0.6, label="Image Overlay Factor"), | |
| gr.Checkbox(True, label="Show Misclassified Images?"), | |
| gr.Slider(value=10, maximum=25, minimum=5, step=5.0, label="#Misclassified images to show"), | |
| gr.Checkbox(True, label="Visulize GradCAM for Misclassified images?"), | |
| gr.Slider(value=10, maximum=25, minimum=5, step=5.0, label="#GradCAM images to show"), | |
| ], | |
| outputs=[ | |
| gr.Label(label="Confidences", container=True, show_label=True), | |
| gr.Image(label="Grad CAM/ Input Image", container=True, show_label=True,height=256,width=256), | |
| gr.Plot(label="Misclassified images", container=True, show_label=True), | |
| gr.Plot(label="Grad CAM of Misclassified images", container=True, show_label=True), | |
| ], | |
| title=TITLE, | |
| description=DESCRIPTION, | |
| examples=examples, | |
| ) | |
| inference_app.launch() | |