Add application file

Bamboo_v0-1_ViT-B16.pth.tar.convert

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

README copy.md

@@ -0,0 +1,13 @@
+---
+title: Bamboo ViT-B16 Demo
+emoji: 🎋
+colorFrom: blue
+colorTo: blue
+sdk: gradio
+sdk_version: 3.0.17
+app_file: app.py
+pinned: false
+license: cc-by-4.0
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,105 @@
+import argparse
+import requests
+import gradio as gr
+import numpy as np
+import cv2
+import torch
+import torch.nn as nn
+from PIL import Image
+import torchvision
+from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
+from timm.data import create_transform
+
+from timmvit import timmvit
+import json
+from timm.models.hub import download_cached_file
+from PIL import Image
+
+def pil_loader(filepath):
+    with Image.open(filepath) as img:
+        img = img.convert('RGB')
+    return img
+
+def build_transforms(input_size, center_crop=True):
+    transform = torchvision.transforms.Compose([
+        torchvision.transforms.ToPILImage(),
+        torchvision.transforms.Resize(input_size * 8 // 7),
+        torchvision.transforms.CenterCrop(input_size),
+        torchvision.transforms.ToTensor(),
+        torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+             ])
+    return transform
+
+# Download human-readable labels for Bamboo.
+with open('./trainid2name.json') as f:
+    id2name = json.load(f)
+
+
+'''
+build model
+'''
+model = timmvit(pretrain_path='./Bamboo_v0-1_ViT-B16.pth.tar.convert')
+model.eval()
+
+'''
+borrow code from here: https://github.com/jacobgil/pytorch-grad-cam/blob/master/pytorch_grad_cam/utils/image.py
+'''
+def show_cam_on_image(img: np.ndarray,
+                      mask: np.ndarray,
+                      use_rgb: bool = False,
+                      colormap: int = cv2.COLORMAP_JET) -> np.ndarray:
+    """ This function overlays the cam mask on the image as an heatmap.
+    By default the heatmap is in BGR format.
+    :param img: The base image in RGB or BGR format.
+    :param mask: The cam mask.
+    :param use_rgb: Whether to use an RGB or BGR heatmap, this should be set to True if 'img' is in RGB format.
+    :param colormap: The OpenCV colormap to be used.
+    :returns: The default image with the cam overlay.
+    """
+    heatmap = cv2.applyColorMap(np.uint8(255 * mask), colormap)
+    if use_rgb:
+        heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
+    heatmap = np.float32(heatmap) / 255
+
+    if np.max(img) > 1:
+        raise Exception(
+            "The input image should np.float32 in the range [0, 1]")
+
+    cam = 0.7*heatmap + 0.3*img
+    # cam = cam / np.max(cam)
+    return np.uint8(255 * cam)
+
+
+
+
+def recognize_image(image):
+    img_t = eval_transforms(image) 
+    # compute output
+    output = model(img_t.unsqueeze(0))
+    prediction = output.softmax(-1).flatten()
+    _,top5_idx = torch.topk(prediction, 5)
+    return {id2name[str(i)][0]: float(prediction[i]) for i in top5_idx.tolist()}
+
+eval_transforms = build_transforms(224)
+
+
+image = gr.inputs.Image()
+label = gr.outputs.Label(num_top_classes=5)
+
+gr.Interface(
+    description="Bamboo for Image Recognition Demo (https://github.com/Davidzhangyuanhan/Bamboo). Bamboo knows what this object is and what you are doing in a very fine-grain granularity: fratercula arctica (fig.5) and dribbler (fig.2)).",
+    fn=recognize_image,
+    inputs=["image"],
+    outputs=[                   
+        label,
+    ],
+    examples=[
+    ["./examples/playing_mahjong.jpg"], 
+    ["./examples/dribbler.jpg"], 
+    ["./examples/Ferrari-F355.jpg"], 
+    ["./examples/northern_oriole.jpg"], 
+    ["./examples/fratercula_arctica.jpg"], 
+    ["./examples/husky.jpg"], 
+    ["./examples/taraxacum_erythrospermum.jpg"], 
+    ],
+).launch()

requirements.txt

@@ -0,0 +1,6 @@
+torchvision==0.11.2
+torch==1.10.1
+opencv-python-headless==4.5.3.56
+timm==0.4.12
+numpy
+

timmvit.py

@@ -0,0 +1,83 @@
+# ------------------------------------------------------------------------
+# SenseTime VTAB
+# Copyright (c) 2021 SenseTime. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+# Modified from DETR (https://github.com/facebookresearch/detr)
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+# ------------------------------------------------------------------------
+
+import timm
+import torch
+import copy
+import torch.nn as nn
+import torchvision
+import json
+from timm.models.hub import download_cached_file
+from PIL import Image
+
+
+
+class MyViT(nn.Module):
+    def __init__(self, num_classes=115217, pretrain_path=None, enable_fc=False):
+        super().__init__()
+        print('initializing ViT model as backbone using ckpt:', pretrain_path)
+        self.model = timm.create_model('vit_base_patch16_224',checkpoint_path=pretrain_path,num_classes=num_classes)# pretrained=True)
+    # def forward_features(self, x):
+    #     x = self.model.patch_embed(x)
+    #     cls_token = self.model.cls_token.expand(x.shape[0], -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+    #     if self.model.dist_token is None:
+    #         x = torch.cat((cls_token, x), dim=1)
+    #     else:
+    #         x = torch.cat((cls_token, self.model.dist_token.expand(x.shape[0], -1, -1), x), dim=1)
+
+    #     x = self.model.pos_drop(x + self.model.pos_embed)
+    #     x = self.model.blocks(x)
+    #     x = self.model.norm(x)
+
+        # return self.model.pre_logits(x[:, 0])
+
+
+    def forward(self, x):
+        x = self.model.forward(x)
+        return x
+
+
+def timmvit(**kwargs):
+    default_kwargs={}
+    default_kwargs.update(**kwargs)
+    return MyViT(**default_kwargs)
+
+
+def build_transforms(input_size, center_crop=True):
+    transform = torchvision.transforms.Compose([
+        torchvision.transforms.Resize(input_size * 8 // 7),
+        torchvision.transforms.CenterCrop(input_size),
+        torchvision.transforms.ToTensor(),
+        torchvision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+             ])
+    return transform
+
+def pil_loader(filepath):
+    with Image.open(filepath) as img:
+        img = img.convert('RGB')
+    return img
+
+def test_build():
+    with open('/mnt/lustre/yhzhang/bamboo/Bamboo_ViT-B16_demo/trainid2name.json') as f:
+        id2name = json.load(f)
+    img = pil_loader('/mnt/lustre/yhzhang/bamboo/Bamboo_ViT-B16_demo/142520422_6ad756ddf6_w_d.jpg')
+    eval_transforms = build_transforms(224)
+    img_t = eval_transforms(img) 
+    img_t = img_t[None, :]
+    model = MyViT(pretrain_path='/mnt/lustre/yhzhang/bamboo/Bamboo_ViT-B16_demo/Bamboo_v0-1_ViT-B16.pth.tar.convert')
+    # image = torch.rand(1, 3, 224, 224)
+    output = model(img_t)
+    # import pdb;pdb.set_trace()
+    prediction = output.softmax(-1).flatten()
+    _,top5_idx = torch.topk(prediction, 5)
+    # import pdb;pdb.set_trace()
+    print({id2name[str(i)][0]: float(prediction[i]) for i in top5_idx.tolist()})
+
+if __name__ == '__main__':
+    test_build()