Spaces:
Runtime error
Runtime error
File size: 4,525 Bytes
eb8427a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 |
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
from torchvision import transforms
from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
from timm.data import create_transform
from config import get_config
from model import build_model
# Download human-readable labels for ImageNet.
response = requests.get("https://git.io/JJkYN")
labels = response.text.split("\n")
def parse_option():
parser = argparse.ArgumentParser('UniCL demo script', add_help=False)
parser.add_argument('--cfg', type=str, default="configs/unicl_swin_base.yaml", metavar="FILE", help='path to config file', )
args, unparsed = parser.parse_known_args()
config = get_config(args)
return args, config
def build_transforms(img_size, center_crop=True):
t = [transforms.ToPILImage()]
if center_crop:
size = int((256 / 224) * img_size)
t.append(
transforms.Resize(size)
)
t.append(
transforms.CenterCrop(img_size)
)
else:
t.append(
transforms.Resize(img_size)
)
t.append(transforms.ToTensor())
t.append(transforms.Normalize(IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD))
return transforms.Compose(t)
def build_transforms4display(img_size, center_crop=True):
t = [transforms.ToPILImage()]
if center_crop:
size = int((256 / 224) * img_size)
t.append(
transforms.Resize(size)
)
t.append(
transforms.CenterCrop(img_size)
)
else:
t.append(
transforms.Resize(img_size)
)
t.append(transforms.ToTensor())
return transforms.Compose(t)
args, config = parse_option()
'''
build model
'''
model = build_model(config)
url = './in21k_yfcc14m_gcc15m_swin_base.pth'
checkpoint = torch.load(url, map_location="cpu")
model.load_state_dict(checkpoint["model"])
model.eval()
'''
build data transform
'''
eval_transforms = build_transforms(224, center_crop=True)
display_transforms = build_transforms4display(224, center_crop=True)
'''
build upsampler
'''
# upsampler = nn.Upsample(scale_factor=16, mode='bilinear')
'''
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, texts):
img_t = eval_transforms(image)
img_d = display_transforms(image).permute(1, 2, 0).numpy()
text_embeddings = model.get_text_embeddings(texts.split(';'))
# compute output
feat_img = model.encode_image(img_t.unsqueeze(0))
output = model.logit_scale.exp() * feat_img @ text_embeddings.t()
prediction = output.softmax(-1).flatten()
return {texts.split(';')[i]: float(prediction[i]) for i in range(len(texts.split(';')))}
image = gr.inputs.Image()
label = gr.outputs.Label(num_top_classes=100)
gr.Interface(
description="UniCL for Zero-shot Image Recognition Demo (https://github.com/microsoft/unicl)",
fn=recognize_image,
inputs=["image", "text"],
outputs=[
label,
],
examples=[
["./elephants.png", "an elephant; an elephant walking in the river; four elephants walking in the river"],
["./apple_with_ipod.jpg", "an ipod; an apple with a write note 'ipod'; an apple"],
["./crowd2.jpg", "a street; a street with a woman walking in the middle; a street with a man walking in the middle"],
],
).launch()
|