Spaces:
Running
Running
| import gradio as gr | |
| import torch | |
| import cv2 | |
| ### CAM explainer code from Intel XAI tools (https://github.com/IntelAI/intel-xai-tools) ### | |
| class XGradCAM: | |
| def __init__(self, model, targetLayer, targetClass, image, dims, device): | |
| # set any frozen layers to trainable | |
| # gradcam cannot be calculated without it | |
| for param in model.parameters(): | |
| if not param.requires_grad: | |
| param.requires_grad = True | |
| self.model = model | |
| self.targetLayer = targetLayer | |
| self.targetClass = targetClass | |
| self.image = image | |
| self.dims = dims | |
| self.device = device | |
| def visualize(self): | |
| from pytorch_grad_cam import XGradCAM, GuidedBackpropReLUModel | |
| from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget | |
| from pytorch_grad_cam.utils.image import show_cam_on_image, deprocess_image, preprocess_image | |
| import torch | |
| import cv2 | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| self.model.eval().to(self.device) | |
| image = cv2.resize(self.image, self.dims) | |
| # convert to rgb if image is grayscale | |
| converted = False | |
| if len(image.shape) == 2: | |
| converted = True | |
| image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB) | |
| rgb_img = np.float32(image) / 255 | |
| input_tensor = preprocess_image(rgb_img, | |
| mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]) | |
| input_tensor = input_tensor.to(self.device) | |
| self.targetLayer = [self.targetLayer] | |
| if self.targetClass is None: | |
| targets = None | |
| else: | |
| targets = [ClassifierOutputTarget(self.targetClass)] | |
| cam = XGradCAM(self.model, self.targetLayer, use_cuda=torch.cuda.is_available()) | |
| # convert back to grayscale if that is the initial dim | |
| if converted: | |
| input_tensor = input_tensor[:, 0:1, :, :] | |
| grayscale_cam = cam(input_tensor=input_tensor, targets=targets, aug_smooth=False, | |
| eigen_smooth=False) | |
| grayscale_cam = grayscale_cam[0, :] | |
| cam_image = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True) | |
| cam_image = cv2.cvtColor(cam_image, cv2.COLOR_RGB2BGR) | |
| gb_model = GuidedBackpropReLUModel(model=self.model, use_cuda=torch.cuda.is_available()) | |
| gb = gb_model(input_tensor, target_category=None) | |
| cam_mask = cv2.merge([grayscale_cam, grayscale_cam, grayscale_cam]) | |
| cam_gb = deprocess_image(cam_mask * gb) | |
| gb = deprocess_image(gb) | |
| print("XGradCAM, Guided backpropagation, and Guided XGradCAM are generated. ") | |
| return cv2.cvtColor(cam_image, cv2.COLOR_RGB2BGR) | |
| class EigenCAM: | |
| def __init__(self, model, targetLayer, boxes, classes, colors, reshape, image, device): | |
| self.model = model | |
| self.targetLayer = targetLayer | |
| self.boxes = boxes | |
| self.classes = classes | |
| self.colors = colors | |
| self.reshape = reshape | |
| self.image = image | |
| self.device = device | |
| def visualize(self): | |
| from pytorch_grad_cam import EigenCAM | |
| from pytorch_grad_cam.utils.image import show_cam_on_image, preprocess_image, scale_cam_image | |
| import torchvision | |
| import torch | |
| import cv2 | |
| import numpy as np | |
| self.model.eval().to(self.device) | |
| rgb_img = np.float32(self.image) / 255 | |
| transform = torchvision.transforms.ToTensor() | |
| input_tensor = transform(rgb_img) | |
| input_tensor = input_tensor.unsqueeze(0) | |
| input_tensor = input_tensor.to(self.device) | |
| self.targetLayer = [self.targetLayer] | |
| if self.reshape is None: | |
| cam = EigenCAM(self.model, self.targetLayer, use_cuda=torch.cuda.is_available()) | |
| else: | |
| cam = EigenCAM(self.model, self.targetLayer, use_cuda=torch.cuda.is_available(), | |
| reshape_transform=self.reshape) | |
| targets = [] | |
| grayscale_cam = cam(input_tensor=input_tensor, targets=targets, aug_smooth=False, | |
| eigen_smooth=False) | |
| grayscale_cam = grayscale_cam[0, :] | |
| cam_image = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True) | |
| renormalized_cam = np.zeros(grayscale_cam.shape, dtype=np.float32) | |
| for x1, y1, x2, y2 in self.boxes: | |
| renormalized_cam[y1:y2, x1:x2] = scale_cam_image(grayscale_cam[y1:y2, x1:x2].copy()) | |
| renormalized_cam = scale_cam_image(renormalized_cam) | |
| eigencam_image_renormalized = show_cam_on_image(rgb_img, renormalized_cam, use_rgb=True) | |
| for i, box in enumerate(self.boxes): | |
| color = self.colors[i] | |
| cv2.rectangle( | |
| eigencam_image_renormalized, | |
| (box[0], box[1]), | |
| (box[2], box[3]), | |
| color, 2 | |
| ) | |
| cv2.putText(eigencam_image_renormalized, self.classes[i], (box[0], box[1] - 5), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2, | |
| lineType=cv2.LINE_AA) | |
| print("EigenCAM is generated. ") | |
| return eigencam_image_renormalized | |
| ### For Gradio Demo ### | |
| def xgradcam(image, model_code, target_class): | |
| global model, target_layer | |
| exec(model_code, globals()) | |
| if target_class == "": | |
| target_class = None | |
| else: | |
| target_class = int(target_class) | |
| image_dims = (224, 224) | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| xgradcam = XGradCAM(model, target_layer, target_class, image, image_dims, device) | |
| return xgradcam.visualize() | |
| def eigencam(image, model_code, class_code, process_code, reshape_code): | |
| global input_image, model, target_layer, bounding_box_coordinates, class_names, box_colors, reshape | |
| input_image = cv2.resize(image, (640, 640)) | |
| exec(model_code, globals()) | |
| exec(class_code, globals()) | |
| exec(process_code, globals()) | |
| exec(reshape_code, globals()) | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| eigencam = EigenCAM(model, target_layer, bounding_box_coordinates, class_names, box_colors, reshape, input_image, device) | |
| return eigencam.visualize() | |
| with gr.Blocks() as demo: | |
| gr.Markdown( | |
| """ | |
| # Class Activation Mapping (CAM) Explainer Demo | |
| This is a demo for CAM explainer from Intel XAI tools (https://github.com/IntelAI/intel-xai-tools). \ | |
| CAM is an approach which localizes regions in the image responsible for a class prediction. \ | |
| The demo shows visualization of XGradCAM for object classification model and EigenCAM for object detection model. | |
| """ | |
| ) | |
| with gr.Tab("XGradCAM"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| xgradcam_image = gr.Image(label="Input Image") | |
| gr.Markdown( | |
| """ | |
| Load the pretrained model to the variable <code>model</code> depending on how it was saved. Then, specify <code>target_layer</code> (normally the last convolutional layer) to compute CAM for. \ | |
| Here are some common choices: | |
| - FasterRCNN: <code>model.backbone</code> | |
| - ResNet18 and 50: <code>model.layer4</code> | |
| - VGG and DenseNet161: <code>model.features</code> | |
| Please don't change the variable names in the following code. | |
| """ | |
| ) | |
| xgradcam_model = gr.Code(label="Model and Target Layer", value= | |
| """ | |
| from torchvision.models import resnet50, ResNet50_Weights | |
| model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2) | |
| target_layer = model.layer4 | |
| """, language="python") | |
| gr.Markdown( | |
| """ | |
| Enter the target category as an integer to compute CAM for. It is the category index in the range <code>[0, NUM_OF_CLASSES-1]</code> based on the training dataset. \ | |
| If it is left blank, the highest scoring category will be used. | |
| """ | |
| ) | |
| xgradcam_targetClass = gr.Textbox(label="Target Category") | |
| xgradcam_output = gr.Image() | |
| xgradcam_button = gr.Button("Submit") | |
| with gr.Tab("EigenCAM"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| eigencam_image = gr.Image(label="Input Image") | |
| gr.Markdown( | |
| """ | |
| Load the pretrained model to the variable <code>model</code> depending on how it was saved. Then, specify <code>target_layer</code> (normally the last convolutional layer) to compute CAM for. \ | |
| Here are some common choices: | |
| - FasterRCNN: <code>model.backbone</code> | |
| - ResNet18 and 50: <code>model.layer4</code> | |
| - VGG and DenseNet161: <code>model.features</code> | |
| Please don't change the variable names in the following code. | |
| """ | |
| ) | |
| eigencam_model = gr.Code(label="Model and Target Layer", value= | |
| """ | |
| from torchvision.models.detection import fasterrcnn_resnet50_fpn | |
| model = fasterrcnn_resnet50_fpn(pretrained=True).eval() | |
| target_layer = model.backbone | |
| """, language="python") | |
| gr.Markdown( | |
| """ | |
| In the case there is no class name in the output from the model, specify <code>class_labels</code> as a list to print them with corresponding bounding box in the image. \ | |
| Depending on the model, the class name might not be needed (e.g. YOLO). Then, create <code>color</code> as a list with a size of the number of classes. | |
| """ | |
| ) | |
| eigencam_class = gr.Code(label="Class Name", value= | |
| """ | |
| import numpy as np | |
| class_labels = ['__background__', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', | |
| 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', | |
| 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', | |
| 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', | |
| 'N/A', 'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', | |
| 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', | |
| 'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass', 'cup', 'fork', | |
| 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', | |
| 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', | |
| 'potted plant', 'bed', 'N/A', 'dining table', 'N/A', 'N/A', 'toilet', | |
| 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', | |
| 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book', 'clock', 'vase', | |
| 'scissors', 'teddy bear', 'hair drier', 'toothbrush'] | |
| color = np.random.uniform(0, 255, size=(len(class_labels), 3)) | |
| """, language="python") | |
| gr.Markdown( | |
| """ | |
| Get <code>output</code> of the model (in the case of FasterRCNN, convert <code>input_image</code> to a tensor first). Then, write a custom <code>process_output</code> function to process the outputs from the model. \ | |
| You should get <code>bounding_box_coordinates</code>, <code>class_names</code>, and <code>box_colors</code> of the detected objects with a higher detection score than <code>detection_threshold</code> value. \ | |
| If you use other models than FasterRCNN, you need to make your own custom process function to match the structure of the outputs from this function. | |
| """ | |
| ) | |
| eigencam_process = gr.Code(label="Output Processing", value= | |
| """ | |
| import torchvision | |
| transform = torchvision.transforms.ToTensor() | |
| input_tensor = transform(np.float32(input_image) / 255).unsqueeze(0) | |
| output = model(input_tensor)[0] | |
| def process_output(output, class_labels, color, detection_threshold): | |
| boxes, classes, labels, colors = [], [], [], [] | |
| box = output['boxes'].tolist() | |
| name = [class_labels[i] for i in output['labels'].detach().numpy()] | |
| label = output['labels'].detach().numpy() | |
| for i in range(len(name)): | |
| score = output['scores'].detach().numpy()[i] | |
| if score < detection_threshold: | |
| continue | |
| boxes.append([int(b) for b in box[i]]) | |
| classes.append(name[i]) | |
| colors.append(color[label[i]]) | |
| return boxes, classes, colors | |
| detection_threshold = 0.9 | |
| bounding_box_coordinates, class_names, box_colors = process_output(output, class_labels, color, detection_threshold) | |
| """, language="python") | |
| gr.Markdown( | |
| """ | |
| Write a custom <code>reshape</code> function to get the activations from the model and process them into 2D format. \ | |
| For example, the backbone of FasterRCNN outputs 5 different tenors with different spatial size as an Ordered Dict, \ | |
| thus, we need a custom function which aggregates these image tensors, resizes them to a common shape, and concatenates them. \ | |
| If you use other models than FasterRCNN, you need to write your own custom reshape function. | |
| """ | |
| ) | |
| eigencam_reshape = gr.Code(label="Reshape", value= | |
| """ | |
| def reshape(x): | |
| target_size = x['pool'].size()[-2 : ] | |
| activations = [] | |
| for key, value in x.items(): | |
| activations.append(torch.nn.functional.interpolate(torch.abs(value), target_size, mode='bilinear')) | |
| activations = torch.cat(activations, axis=1) | |
| return activations | |
| """, language="python") | |
| eigencam_output = gr.Image() | |
| eigencam_button = gr.Button("Submit") | |
| xgradcam_button.click(xgradcam, inputs=[xgradcam_image, xgradcam_model, xgradcam_targetClass], outputs=xgradcam_output) | |
| eigencam_button.click(eigencam, inputs=[eigencam_image, eigencam_model, eigencam_class, eigencam_process, eigencam_reshape], outputs=eigencam_output) | |
| demo.launch() | |