import copy import os import sys sys.path.append('src') import shutil from collections import defaultdict from functools import lru_cache import cv2 import gradio as gr import mediapy import numpy as np import pandas as pd import torch from deep_translator import GoogleTranslator from gradio_blocks import build_video_to_camvideo from Nets import CustomResNet18 from PIL import Image from pytorch_grad_cam import GradCAM, HiResCAM, GradCAMPlusPlus, AblationCAM, XGradCAM, EigenCAM, FullGrad from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget from pytorch_grad_cam.utils.image import show_cam_on_image from tqdm import tqdm import util from util import transform, CustomImageCache, imageCacheWrapper util.ImageCache = CustomImageCache(60, False) ffmpeg_path = shutil.which('ffmpeg') mediapy.set_ffmpeg(ffmpeg_path) IMAGE_PATH = os.path.join(os.getcwd(), 'src/examples') IMAGES_PER_ROW = 5 MAXIMAL_FRAMES = 1000 BATCHES_TO_PROCESS = 10 OUTPUT_FPS = 10 MAX_OUT_FRAMES = 60 MODEL = CustomResNet18(90).eval() MODEL.load_state_dict(torch.load('src/results/models/best_model.pth', map_location=torch.device('cpu'))) CAM_METHODS = { "GradCAM": GradCAM, "GradCAM++": GradCAMPlusPlus, "XGradCAM": XGradCAM, "HiResCAM": HiResCAM, "EigenCAM": EigenCAM } LAYERS = { 'layer1': MODEL.resnet.layer1, 'layer2': MODEL.resnet.layer2, 'layer3': MODEL.resnet.layer3, 'layer4': MODEL.resnet.layer4, 'all': [MODEL.resnet.layer1, MODEL.resnet.layer2, MODEL.resnet.layer3, MODEL.resnet.layer4], 'layer3+4': [MODEL.resnet.layer3, MODEL.resnet.layer4] } CV2_COLORMAPS = { "Autumn": cv2.COLORMAP_AUTUMN, "Bone": cv2.COLORMAP_BONE, "Jet": cv2.COLORMAP_JET, "Winter": cv2.COLORMAP_WINTER, "Rainbow": cv2.COLORMAP_RAINBOW, "Ocean": cv2.COLORMAP_OCEAN, "Summer": cv2.COLORMAP_SUMMER, "Pink": cv2.COLORMAP_PINK, "Hot": cv2.COLORMAP_HOT, "Magma": cv2.COLORMAP_MAGMA, "Inferno": cv2.COLORMAP_INFERNO, "Plasma": cv2.COLORMAP_PLASMA, "Twilight": cv2.COLORMAP_TWILIGHT, } # cam_model = copy.deepcopy(model) data_df = pd.read_csv('src/cache/val_df.csv') C_NUM_TO_NAME = data_df[['encoded_target', 'target']].drop_duplicates().sort_values('encoded_target').set_index('encoded_target')['target'].to_dict() C_NAME_TO_NUM = {v: k for k, v in C_NUM_TO_NAME.items()} ALL_CLASSES = sorted(list(C_NUM_TO_NAME.values()), key=lambda x: x.lower()) def get_class_name(idx): return C_NUM_TO_NAME[idx] def get_class_idx(name): return C_NAME_TO_NUM[name] @lru_cache(maxsize=100) def get_translated(to_translate): return GoogleTranslator(source="en", target="de").translate(to_translate) for idx in range(90): get_translated(get_class_name(idx)) @imageCacheWrapper def infer_image(image): if isinstance(image, dict): # Its the image and a mask as pillow both -> Combine them to one image image = Image.blend(image["image"], image["mask"], alpha=0.5) image.save('src/results/infer_image.png') image = transform(image) image = image.unsqueeze(0) with torch.no_grad(): output = MODEL(image) distribution = torch.nn.functional.softmax(output, dim=1) ret = defaultdict(float) for idx, prob in enumerate(distribution[0]): animal = f'{get_class_name(idx)} ({get_translated(get_class_name(idx))})' ret[animal] = prob.item() return ret def gradcam(image, colormap="Jet", use_eigen_smooth=False, use_aug_smooth=False, BWHighlight=False, alpha=0.5, cam_method=GradCAM, layer=None, specific_class="Predicted Class"): if image is None: raise gr.Error("Please upload an image.") if isinstance(image, dict): # Its the image and a mask as pillow both -> Combine them to one image image = Image.blend(image["image"], image["mask"], alpha=0.5) if colormap not in CV2_COLORMAPS.keys(): raise gr.Error(f"Colormap {colormap} not found in {list(CV2_COLORMAPS.keys())}.") else: colormap = CV2_COLORMAPS[colormap] image_width, image_height = image.size if image_width > 4000 or image_height > 4000: raise gr.Error("The image is too big. The maximal size is 4000x4000.") MODEL.eval() layers = LAYERS[layer] image_tensor = transform(image).unsqueeze(0) targets = [ClassifierOutputTarget(get_class_idx(specific_class))] if specific_class != "Predicted Class" else None with CAM_METHODS[cam_method](model=MODEL, target_layers=layers) as cam: grayscale_cam = cam(input_tensor=image_tensor, targets=targets, aug_smooth=use_aug_smooth, eigen_smooth=use_eigen_smooth) grayscale_cam = grayscale_cam[0, :] grayscale_cam = cv2.resize(grayscale_cam, (image_width, image_height), interpolation=cv2.INTER_CUBIC) image = np.float32(image) visualization = None if BWHighlight: image = image * grayscale_cam[..., np.newaxis] visualization = image.astype(np.uint8) else: image = image / 255 visualization = show_cam_on_image(image, grayscale_cam, use_rgb=True, image_weight=alpha, colormap=colormap) return Image.fromarray(visualization) def gradcam_video(video, colormap="Jet", use_eigen_smooth=False, BWHighlight=False, alpha=0.5, cam_method=GradCAM, layer=None, specific_class="Predicted Class"): global OUTPUT_FPS, MAXIMAL_FRAMES, BATCHES_TO_PROCESS, MAX_OUT_FRAMES if colormap not in CV2_COLORMAPS.keys(): raise gr.Error(f"Colormap {colormap} not found in {list(CV2_COLORMAPS.keys())}.") else: colormap = CV2_COLORMAPS[colormap] video = cv2.VideoCapture(video) fps = int(video.get(cv2.CAP_PROP_FPS)) if OUTPUT_FPS == -1: OUTPUT_FPS = fps width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)) if width > 3000 or height > 3000: raise gr.Error("The video is too big. The maximal size is 3000x3000.") print(f'FPS: {fps}, Width: {width}, Height: {height}') frames = list() success, image = video.read() while success: frames.append(image) success, image = video.read() print(f'Frames: {len(frames)}') if len(frames) == 0: raise gr.Error("The video is empty.") if len(frames) >= MAXIMAL_FRAMES: raise gr.Error(f"The video is too long. The maximal length is {MAXIMAL_FRAMES} frames.") if len(frames) > MAX_OUT_FRAMES: frames = frames[::len(frames) // MAX_OUT_FRAMES] print(f'Frames to process: {len(frames)}') processed = [Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) for frame in frames] # generate lists in lists for the images for batch processing. BATCHES_TO_PROCESS images per inner list batched = [processed[i:i + BATCHES_TO_PROCESS] for i in range(0, len(processed), BATCHES_TO_PROCESS)] MODEL.eval() layers = LAYERS[layer] results = list() targets = [ClassifierOutputTarget(get_class_idx(specific_class))] if specific_class != "Predicted Class" else None with CAM_METHODS[cam_method](model=MODEL, target_layers=layers) as cam: for i, batch in enumerate(tqdm(batched)): images_tensor = torch.stack([transform(image) for image in batch]) grayscale_cam = cam(input_tensor=images_tensor, targets=targets, aug_smooth=False, eigen_smooth=use_eigen_smooth) for i, image in enumerate(batch): _grayscale_cam = grayscale_cam[i, :] _grayscale_cam = cv2.resize(_grayscale_cam, (width, height), interpolation=cv2.INTER_LINEAR) image = np.float32(image) visualization = None if BWHighlight: image = image * _grayscale_cam[..., np.newaxis] visualization = image.astype(np.uint8) else: image = image / 255 visualization = show_cam_on_image(image, _grayscale_cam, use_rgb=True, image_weight=alpha, colormap=colormap) results.append(visualization) # save video mediapy.write_video('src/results/gradcam_video.mp4', results, fps=OUTPUT_FPS) video.release() return 'src/results/gradcam_video.mp4' def load_examples(): folder_name_to_header = { "AI_Generated": "AI Generated Images", "true_predicted": "True Predicted Images (Validation Set)", "false_predicted": "False Predicted Images (Validation Set)", "others": "Other interesting images from the internet" } images_description = { "AI_Generated": "These images are generated by Dalle3 and Stable Diffusion. All of them are not real images and because of that it is interesting to see how the model predicts them.", "true_predicted": "These images are from the validation set and the model predicted them correctly.", "false_predicted": "These images are from the validation set and the model predicted them incorrectly. Maybe you can see why the model predicted them incorrectly using the GradCAM visualization. :)", "others": "These images are from the internet and are not part of the validation set. They are interesting because most of them show different animals." } loaded_images = defaultdict(list) for image_type in ["AI_Generated", "true_predicted", "false_predicted", "others"]: # for image_type in os.listdir(IMAGE_PATH): full_path = os.path.join(IMAGE_PATH, image_type).replace('\\', '/').replace('//', '/') gr.Markdown(f'## {folder_name_to_header[image_type]}') gr.Markdown(images_description[image_type]) images_to_load = os.listdir(full_path) rows = (len(images_to_load) // IMAGES_PER_ROW) + 1 for i in range(rows): with gr.Row(elem_classes=["row-example-images"], equal_height=False): for j in range(IMAGES_PER_ROW): if i * IMAGES_PER_ROW + j >= len(images_to_load): break image = images_to_load[i * IMAGES_PER_ROW + j] name = f"{image.split('.')[0]} ({get_translated(image.split('.')[0])})" image = Image.open(os.path.join(full_path, image)) # scale so that the longest side is 600px scale = 600 / max(image.size) image = image.resize((int(image.size[0] * scale), int(image.size[1] * scale))) loaded_images[image_type].append( gr.Image( value=image, label=name, type="pil", interactive=False, elem_classes=["selectable_images"], ) ) return loaded_images css = """ #logo {text-align: right;} p {text-align: justify; text-justify: inter-word; font-size: 1.1em; line-height: 1.2em;} .svelte-1btp92j.selectable {cursor: pointer !important; } """ with gr.Blocks(theme='freddyaboulton/dracula_revamped', css=css) as demo: # ------------------------------------------- # HEADER WITH LOGO # ------------------------------------------- with gr.Row(): with open('src/header.md', 'r', encoding='utf-8') as f: markdown_string = f.read() with gr.Column(scale=10): header = gr.Markdown(markdown_string) with gr.Column(scale=1): pil_logo = Image.open('animals.png') logo = gr.Image(value=pil_logo, scale=2, interactive=False, show_download_button=False, show_label=False, container=False, elem_id="logo") # ------------------------------------------- # INPUT IMAGE # ------------------------------------------- with gr.Row(): with gr.Row(variant="panel", equal_height=True): user_image = gr.Image( type="pil", label="Upload Your Own Image", tool="sketch", interactive=True, ) # ------------------------------------------- # TOOLS # ------------------------------------------- with gr.Row(): # ------------------------------------------- # PREDICT # ------------------------------------------- with gr.Tab("Predict"): with gr.Column(): output = gr.Label( num_top_classes=5, label="Output", info="Top three predicted classes and their confidences.", scale=5, ) predict_mode_button = gr.Button(value="Predict Animal", label="Predict", info="Click to make a prediction.", scale=1) predict_mode_button.click(fn=infer_image, inputs=[user_image], outputs=output, queue=True) # ------------------------------------------- # EXPLAIN # ------------------------------------------- with gr.Tab("Explain Image"): with gr.Row(): with gr.Column(): _info = "There are different GradCAM methods. You can read more about them here: (https://github.com/jacobgil/pytorch-grad-cam#references)." cam_method = gr.Radio( list(CAM_METHODS.keys()), label="GradCAM Method", info=_info, value="GradCAM", interactive=True, scale=2, ) _info = """ The alpha value is used to blend the original image with the GradCAM visualization. If you choose a value of 0.5 the original image and the GradCAM visualization will be blended equally. If you choose a value of 0.1 the original image will be barely visible and if you choose a value of 0.9 the GradCAM visualization will be barely visible. """ alpha = gr.Slider( minimum=.1, maximum=.9, value=0.5, interactive=True, step=.1, label="Alpha", scale=1, info=_info ) _info = """ The layer is used to choose the layer of the ResNet50 model. The GradCAM visualization will be based on this layer. Best to choose is the last layer (layer4) because it is the layer with the most information before the final prediction. This makes the GradCAM visualization the most meaningful. If all layers are chosen the GradCAM visualization will be averaged over all layers. """ layer = gr.Radio( LAYERS.keys(), label="Layer", value="layer4", interactive=True, scale=2, info=_info ) _info = """ Here you can choose the animal to "explain". If you choose "Predicted Class" the GradCAM visualization will be based on the predicted class. If you choose a specific class the GradCAM visualization will be based on this class. For example if you have an image with a dog and a cat, you can select either Cat or Dog and see if the model can focus on the correct animal. """ animal_to_explain = gr.Dropdown( choices=["Predicted Class"] + ALL_CLASSES, label="Animal", value="Predicted Class", interactive=True, scale=2, info=_info ) with gr.Row(): _info = """ Here you can choose the colormap. Instead of a colormap you can also choose "BW Highlight" to just keep the original image and highlight the important parts of the image. If you select "BW Highlight" the colormap will be ignored. """ colormap = gr.Dropdown( choices=list(CV2_COLORMAPS.keys()), label="Colormap", value="Jet", interactive=True, scale=2, info=_info ) bw_highlight = gr.Checkbox( label="BW Highlight", value=False, interactive=True, scale=1, ) bw_highlight.description = "Here you can choose if you want to highlight the important parts of the image in black and white." with gr.Row(): _info = """ The Eigen Smooth is a method to smooth the GradCAM visualization. """ use_eigen_smooth = gr.Checkbox( label="Eigen Smooth", value=False, interactive=True, scale=1, info=_info ) _info = """ The Aug Smooth is also a method to smooth the GradCAM visualization. But this method needs a lot of performance and is therefore slow. """ use_aug_smooth = gr.Checkbox( label="Aug Smooth", value=False, interactive=True, scale=1, info=_info ) with gr.Column(): output_cam = gr.Image( type="pil", label="GradCAM", info="GradCAM visualization", scale=5, ) gradcam_mode_button = gr.Button(value="Show GradCAM", label="GradCAM", info="Click to make a prediction.", scale=1) gradcam_mode_button.click(fn=gradcam, inputs=[user_image, colormap, use_eigen_smooth, use_aug_smooth, bw_highlight, alpha, cam_method, layer, animal_to_explain], outputs=output_cam, queue=True) # ------------------------------------------- # Video CAM # ------------------------------------------- with gr.Tab("Explain Video"): build_video_to_camvideo(CAM_METHODS, CV2_COLORMAPS, LAYERS, ALL_CLASSES, gradcam_video) # ------------------------------------------- # EXAMPLES # ------------------------------------------- with gr.Tab("Example Images"): placeholder = gr.Markdown("## Example Images") loaded_images = load_examples() for k in loaded_images.keys(): for image in loaded_images[k]: image.select(fn=lambda x: x, inputs=[image], outputs=[user_image]) if __name__ == "__main__": demo.queue() demo.launch()