initial commit

app.py

@@ -0,0 +1,310 @@
+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()

model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d04e3718649a13a5a49ae5c274ae3aefb9deb229ef5194106bb8ecbfd4f00c61
+size 704238867

requirements.txt

@@ -0,0 +1,7 @@
+torch
+torchvision
+pillow
+numpy
+scipy
+einops
+opencv-python-headless

vision_transformer.py

@@ -0,0 +1,330 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# 
+#     http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Mostly copy-paste from timm library.
+https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
+"""
+import math
+from functools import partial
+
+import torch
+import torch.nn as nn
+
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+                      "The distribution of values may be incorrect.",
+                      stacklevel=2)
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        l = norm_cdf((a - mean) / std)
+        u = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [l, u], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    # type: (Tensor, float, float, float, float) -> Tensor
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+
+def drop_path(x, drop_prob: float = 0., training: bool = False):
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()  # binarize
+    output = x.div(keep_prob) * random_tensor
+    return output
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x, attn
+
+
+class Block(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+    def forward(self, x, return_attention=False):
+        y, attn = self.attn(self.norm1(x))
+        if return_attention:
+            return attn
+        x = x + self.drop_path(y)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        num_patches = (img_size // patch_size) * (img_size // patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x
+
+
+class VisionTransformer(nn.Module):
+    """ Vision Transformer """
+    def __init__(self, img_size=[224], patch_size=16, in_chans=3, num_classes=0, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
+                 drop_path_rate=0., norm_layer=nn.LayerNorm, **kwargs):
+        super().__init__()
+        self.num_features = self.embed_dim = embed_dim
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size[0], patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer)
+            for i in range(depth)])
+        self.norm = norm_layer(embed_dim)
+
+        # Classifier head
+        self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+        trunc_normal_(self.pos_embed, std=.02)
+        trunc_normal_(self.cls_token, std=.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def interpolate_pos_encoding(self, x, w, h):
+        npatch = x.shape[1] - 1
+        N = self.pos_embed.shape[1] - 1
+        if npatch == N and w == h:
+            return self.pos_embed
+        class_pos_embed = self.pos_embed[:, 0]
+        patch_pos_embed = self.pos_embed[:, 1:]
+        dim = x.shape[-1]
+        w0 = w // self.patch_embed.patch_size
+        h0 = h // self.patch_embed.patch_size
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        w0, h0 = w0 + 0.1, h0 + 0.1
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
+            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
+            mode='bicubic',
+        )
+        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
+
+    def prepare_tokens(self, x):
+        B, nc, w, h = x.shape
+        x = self.patch_embed(x)  # patch linear embedding
+
+        # add the [CLS] token to the embed patch tokens
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        # add positional encoding to each token
+        x = x + self.interpolate_pos_encoding(x, w, h)
+
+        return self.pos_drop(x)
+
+    def forward(self, x):
+        x = self.prepare_tokens(x)
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+        return x[:, 0]
+
+    def get_last_selfattention(self, x):
+        x = self.prepare_tokens(x)
+        for i, blk in enumerate(self.blocks):
+            if i < len(self.blocks) - 1:
+                x = blk(x)
+            else:
+                # return attention of the last block
+                return blk(x, return_attention=True)
+
+    def get_intermediate_layers(self, x, n=1):
+        x = self.prepare_tokens(x)
+        # we return the output tokens from the `n` last blocks
+        output = []
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if len(self.blocks) - i <= n:
+                output.append(self.norm(x))
+        return output
+
+
+def vit_tiny(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=192, depth=12, num_heads=3, mlp_ratio=4,
+        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def vit_small(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=384, depth=12, num_heads=6, mlp_ratio=4,
+        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+def vit_base(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4,
+        qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+
+class DINOHead(nn.Module):
+    def __init__(self, in_dim, out_dim, use_bn=False, norm_last_layer=True, nlayers=3, hidden_dim=2048, bottleneck_dim=256):
+        super().__init__()
+        nlayers = max(nlayers, 1)
+        if nlayers == 1:
+            self.mlp = nn.Linear(in_dim, bottleneck_dim)
+        else:
+            layers = [nn.Linear(in_dim, hidden_dim)]
+            if use_bn:
+                layers.append(nn.BatchNorm1d(hidden_dim))
+            layers.append(nn.GELU())
+            for _ in range(nlayers - 2):
+                layers.append(nn.Linear(hidden_dim, hidden_dim))
+                if use_bn:
+                    layers.append(nn.BatchNorm1d(hidden_dim))
+                layers.append(nn.GELU())
+            layers.append(nn.Linear(hidden_dim, bottleneck_dim))
+            self.mlp = nn.Sequential(*layers)
+        self.apply(self._init_weights)
+        self.last_layer = nn.utils.weight_norm(nn.Linear(bottleneck_dim, out_dim, bias=False))
+        self.last_layer.weight_g.data.fill_(1)
+        if norm_last_layer:
+            self.last_layer.weight_g.requires_grad = False
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        x = self.mlp(x)
+        x = nn.functional.normalize(x, dim=-1, p=2)
+        x = self.last_layer(x)
+        return x