app.py

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()