Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import os | |
| import sys | |
| import cv2 | |
| import matplotlib | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| from PIL import Image | |
| from PIL import ImageFont | |
| from PIL import ImageDraw | |
| from scipy.stats import rankdata | |
| import torch | |
| import torch.nn as nn | |
| import torchvision | |
| from torchvision import transforms as pth_transforms | |
| import torchvision.transforms as transforms | |
| from einops import rearrange, repeat | |
| import vision_transformer as vits | |
| def get_vit256(pretrained_weights, arch='vit_small', device=torch.device('cpu')): | |
| r""" | |
| Builds ViT-256 Model. | |
| Args: | |
| - pretrained_weights (str): Path to ViT-256 Model Checkpoint. | |
| - arch (str): Which model architecture. | |
| - device (torch): Torch device to save model. | |
| Returns: | |
| - model256 (torch.nn): Initialized model. | |
| """ | |
| checkpoint_key = 'teacher' | |
| device = torch.device("cpu") if torch.cuda.is_available() else torch.device("cpu") | |
| model256 = vits.__dict__[arch](patch_size=16, num_classes=0) | |
| for p in model256.parameters(): | |
| p.requires_grad = False | |
| model256.eval() | |
| model256.to(device) | |
| if os.path.isfile(pretrained_weights): | |
| state_dict = torch.load(pretrained_weights, map_location="cpu") | |
| if checkpoint_key is not None and checkpoint_key in state_dict: | |
| print(f"Take key {checkpoint_key} in provided checkpoint dict") | |
| state_dict = state_dict[checkpoint_key] | |
| # remove `module.` prefix | |
| state_dict = {k.replace("module.", ""): v for k, v in state_dict.items()} | |
| # remove `backbone.` prefix induced by multicrop wrapper | |
| state_dict = {k.replace("backbone.", ""): v for k, v in state_dict.items()} | |
| msg = model256.load_state_dict(state_dict, strict=False) | |
| print('Pretrained weights found at {} and loaded with msg: {}'.format(pretrained_weights, msg)) | |
| return model256 | |
| def cmap_map(function, cmap): | |
| r""" | |
| Applies function (which should operate on vectors of shape 3: [r, g, b]), on colormap cmap. | |
| This routine will break any discontinuous points in a colormap. | |
| Args: | |
| - function (function) | |
| - cmap (matplotlib.colormap) | |
| Returns: | |
| - matplotlib.colormap | |
| """ | |
| cdict = cmap._segmentdata | |
| step_dict = {} | |
| # Firt get the list of points where the segments start or end | |
| for key in ('red', 'green', 'blue'): | |
| step_dict[key] = list(map(lambda x: x[0], cdict[key])) | |
| step_list = sum(step_dict.values(), []) | |
| step_list = np.array(list(set(step_list))) | |
| # Then compute the LUT, and apply the function to the LUT | |
| reduced_cmap = lambda step : np.array(cmap(step)[0:3]) | |
| old_LUT = np.array(list(map(reduced_cmap, step_list))) | |
| new_LUT = np.array(list(map(function, old_LUT))) | |
| # Now try to make a minimal segment definition of the new LUT | |
| cdict = {} | |
| for i, key in enumerate(['red','green','blue']): | |
| this_cdict = {} | |
| for j, step in enumerate(step_list): | |
| if step in step_dict[key]: | |
| this_cdict[step] = new_LUT[j, i] | |
| elif new_LUT[j,i] != old_LUT[j, i]: | |
| this_cdict[step] = new_LUT[j, i] | |
| colorvector = list(map(lambda x: x + (x[1], ), this_cdict.items())) | |
| colorvector.sort() | |
| cdict[key] = colorvector | |
| return matplotlib.colors.LinearSegmentedColormap('colormap',cdict,1024) | |
| def identity(x): | |
| r""" | |
| Identity Function. | |
| Args: | |
| - x: | |
| Returns: | |
| - x | |
| """ | |
| return x | |
| def tensorbatch2im(input_image, imtype=np.uint8): | |
| r"""" | |
| Converts a Tensor array into a numpy image array. | |
| Args: | |
| - input_image (torch.Tensor): (B, C, W, H) Torch Tensor. | |
| - imtype (type): the desired type of the converted numpy array | |
| Returns: | |
| - image_numpy (np.array): (B, W, H, C) Numpy Array. | |
| """ | |
| if not isinstance(input_image, np.ndarray): | |
| image_numpy = input_image.cpu().float().numpy() # convert it into a numpy array | |
| #if image_numpy.shape[0] == 1: # grayscale to RGB | |
| # image_numpy = np.tile(image_numpy, (3, 1, 1)) | |
| image_numpy = (np.transpose(image_numpy, (0, 2, 3, 1)) + 1) / 2.0 * 255.0 # post-processing: tranpose and scaling | |
| else: # if it is a numpy array, do nothing | |
| image_numpy = input_image | |
| return image_numpy.astype(imtype) | |
| def getConcatImage(imgs, how='horizontal', gap=0): | |
| r""" | |
| Function to concatenate list of images (vertical or horizontal). | |
| Args: | |
| - imgs (list of PIL.Image): List of PIL Images to concatenate. | |
| - how (str): How the images are concatenated (either 'horizontal' or 'vertical') | |
| - gap (int): Gap (in px) between images | |
| Return: | |
| - dst (PIL.Image): Concatenated image result. | |
| """ | |
| gap_dist = (len(imgs)-1)*gap | |
| if how == 'vertical': | |
| w, h = np.max([img.width for img in imgs]), np.sum([img.height for img in imgs]) | |
| h += gap_dist | |
| curr_h = 0 | |
| dst = Image.new('RGBA', (w, h), color=(255, 255, 255, 0)) | |
| for img in imgs: | |
| dst.paste(img, (0, curr_h)) | |
| curr_h += img.height + gap | |
| elif how == 'horizontal': | |
| w, h = np.sum([img.width for img in imgs]), np.min([img.height for img in imgs]) | |
| w += gap_dist | |
| curr_w = 0 | |
| dst = Image.new('RGBA', (w, h), color=(255, 255, 255, 0)) | |
| for idx, img in enumerate(imgs): | |
| dst.paste(img, (curr_w, 0)) | |
| curr_w += img.width + gap | |
| return dst | |
| def add_margin(pil_img, top, right, bottom, left, color): | |
| r""" | |
| Adds custom margin to PIL.Image. | |
| """ | |
| width, height = pil_img.size | |
| new_width = width + right + left | |
| new_height = height + top + bottom | |
| result = Image.new(pil_img.mode, (new_width, new_height), color) | |
| result.paste(pil_img, (left, top)) | |
| return result | |
| def concat_scores256(attns, size=(256,256)): | |
| r""" | |
| """ | |
| rank = lambda v: rankdata(v)*100/len(v) | |
| color_block = [rank(attn.flatten()).reshape(size) for attn in attns] | |
| color_hm = np.concatenate([ | |
| np.concatenate(color_block[i:(i+16)], axis=1) | |
| for i in range(0,256,16) | |
| ]) | |
| return color_hm | |
| def get_scores256(attns, size=(256,256)): | |
| r""" | |
| """ | |
| rank = lambda v: rankdata(v)*100/len(v) | |
| color_block = [rank(attn.flatten()).reshape(size) for attn in attns][0] | |
| return color_block | |
| def get_patch_attention_scores(patch, model256, scale=1, device256=torch.device('cpu')): | |
| t = transforms.Compose([ | |
| transforms.ToTensor(), | |
| transforms.Normalize( | |
| [0.5, 0.5, 0.5], [0.5, 0.5, 0.5] | |
| ) | |
| ]) | |
| with torch.no_grad(): | |
| batch_256 = t(patch).unsqueeze(0) | |
| batch_256 = batch_256.to(device256, non_blocking=True) | |
| features_256 = model256(batch_256) | |
| attention_256 = model256.get_last_selfattention(batch_256) | |
| nh = attention_256.shape[1] # number of head | |
| attention_256 = attention_256[:, :, 0, 1:].reshape(256, nh, -1) | |
| attention_256 = attention_256.reshape(1, nh, 16, 16) | |
| attention_256 = nn.functional.interpolate(attention_256, scale_factor=int(16/scale), mode="nearest").cpu().numpy() | |
| if scale != 1: | |
| batch_256 = nn.functional.interpolate(batch_256, scale_factor=(1/scale), mode="nearest") | |
| return tensorbatch2im(batch_256), attention_256 | |
| def create_patch_heatmaps_concat(patch, model256, output_dir=None, fname=None, threshold=None, | |
| offset=16, alpha=0.5, cmap=plt.get_cmap('coolwarm')): | |
| r""" | |
| Creates patch heatmaps (concatenated for easy comparison) | |
| Args: | |
| - patch (PIL.Image): 256 x 256 Image | |
| - model256 (torch.nn): 256-Level ViT | |
| - output_dir (str): Save directory / subdirectory | |
| - fname (str): Naming structure of files | |
| - offset (int): How much to offset (from top-left corner with zero-padding) the region by for blending | |
| - alpha (float): Image blending factor for cv2.addWeighted | |
| - cmap (matplotlib.pyplot): Colormap for creating heatmaps | |
| Returns: | |
| - None | |
| """ | |
| patch1 = patch.copy() | |
| patch2 = add_margin(patch.crop((16,16,256,256)), top=0, left=0, bottom=16, right=16, color=(255,255,255)) | |
| b256_1, a256_1 = get_patch_attention_scores(patch1, model256) | |
| b256_1, a256_2 = get_patch_attention_scores(patch2, model256) | |
| save_region = np.array(patch.copy()) | |
| s = 256 | |
| offset_2 = offset | |
| if threshold != None: | |
| ths = [] | |
| for i in range(6): | |
| score256_1 = get_scores256(a256_1[:,i,:,:], size=(s,)*2) | |
| score256_2 = get_scores256(a256_2[:,i,:,:], size=(s,)*2) | |
| new_score256_2 = np.zeros_like(score256_2) | |
| new_score256_2[offset_2:s, offset_2:s] = score256_2[:(s-offset_2), :(s-offset_2)] | |
| overlay256 = np.ones_like(score256_2)*100 | |
| overlay256[offset_2:s, offset_2:s] += 100 | |
| score256 = (score256_1+new_score256_2)/overlay256 | |
| mask256 = score256.copy() | |
| mask256[mask256 < threshold] = 0 | |
| mask256[mask256 > threshold] = 0.95 | |
| color_block256 = (cmap(mask256)*255)[:,:,:3].astype(np.uint8) | |
| region256_hm = cv2.addWeighted(color_block256, alpha, save_region.copy(), 1-alpha, 0, save_region.copy()) | |
| region256_hm[mask256==0] = 0 | |
| img_inverse = save_region.copy() | |
| img_inverse[mask256 == 0.95] = 0 | |
| ths.append(region256_hm+img_inverse) | |
| ths = [Image.fromarray(img) for img in ths] | |
| getConcatImage([getConcatImage(ths[0:3]), | |
| getConcatImage(ths[4:6])], how='vertical').save(os.path.join(output_dir, '%s_256th.png' % (fname))) | |
| hms = [] | |
| for i in range(6): | |
| score256_1 = get_scores256(a256_1[:,i,:,:], size=(s,)*2) | |
| score256_2 = get_scores256(a256_2[:,i,:,:], size=(s,)*2) | |
| new_score256_2 = np.zeros_like(score256_2) | |
| new_score256_2[offset_2:s, offset_2:s] = score256_2[:(s-offset_2), :(s-offset_2)] | |
| overlay256 = np.ones_like(score256_2)*100 | |
| overlay256[offset_2:s, offset_2:s] += 100 | |
| score256 = (score256_1+new_score256_2)/overlay256 | |
| color_block256 = (cmap(score256)*255)[:,:,:3].astype(np.uint8) | |
| region256_hm = cv2.addWeighted(color_block256, alpha, save_region.copy(), 1-alpha, 0, save_region.copy()) | |
| hms.append(region256_hm) | |
| hms = [Image.fromarray(img) for img in hms] | |
| return getConcatImage([getConcatImage(hms[0:3], how='horizontal', gap=10), | |
| getConcatImage(hms[4:6], how='horizontal', gap=10)], how='vertical', gap=10) | |
| def demo_patch_heatmaps(input_image): | |
| light_jet = cmap_map(lambda x: x/2 + 0.5, matplotlib.cm.jet) | |
| model256 = get_vit256(pretrained_weights=pretrained_weights256) | |
| demo_heatmap = create_patch_heatmaps_concat(input_image, model256, cmap=light_jet) | |
| return demo_heatmap | |
| pretrained_weights256 = './model.pt' | |
| title = "Demo for 11604" | |
| description = "To use, upload a 256 x 256 patch (20X magnification). \ | |
| The output will generate attention results from 6 attention heads." | |
| iface = gr.Interface(fn=demo_patch_heatmaps, | |
| inputs=gr.inputs.Image(type='pil'), | |
| outputs="image", | |
| title=title, | |
| description=description, | |
| allow_flagging=False) | |
| iface.launch() | |