Spaces:
Runtime error
Runtime error
File size: 2,376 Bytes
d4ab5ac 8fd2935 d4ab5ac 8fd2935 d4ab5ac d85fbeb 8fd2935 d4ab5ac 8fd2935 d85fbeb 8fd2935 d4ab5ac d85fbeb d4ab5ac 8fd2935 d85fbeb d4ab5ac 8fd2935 |
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 |
import sys
sys.path.insert(0, './code')
from datamodules.transformations import UnNest
from models.interpretation import ImageInterpretationNet
from transformers import ViTFeatureExtractor, ViTForImageClassification
from utils.plot import smoothen, draw_mask_on_image, draw_heatmap_on_image
import gradio as gr
import numpy as np
import torch
# Load Vision Transformer
hf_model = "tanlq/vit-base-patch16-224-in21k-finetuned-cifar10"
vit = ViTForImageClassification.from_pretrained(hf_model)
vit.eval()
# Load Feature Extractor
feature_extractor = ViTFeatureExtractor.from_pretrained(hf_model, return_tensors="pt")
feature_extractor = UnNest(feature_extractor)
# Load Vision DiffMask
diffmask = ImageInterpretationNet.load_from_checkpoint('checkpoints/diffmask.ckpt')
diffmask.set_vision_transformer(vit)
diffmask_imagenet = ImageInterpretationNet.load_from_checkpoint('checkpoints/diffmask_imagenet.ckpt')
diffmask_imagenet.set_vision_transformer(vit)
# Define mask plotting functions
def draw_mask(image, mask):
return draw_mask_on_image(image, smoothen(mask))\
.permute(1, 2, 0)\
.clip(0, 1)\
.numpy()
def draw_heatmap(image, mask):
return draw_heatmap_on_image(image, smoothen(mask))\
.permute(1, 2, 0)\
.clip(0, 1)\
.numpy()
# Define callable method for the demo
def get_mask(image, model_name: str):
if image is None:
return None, None
if model_name == 'DiffMask-CiFAR-10':
diffmask_model = diffmask
elif model_name == 'DiffMask-ImageNet':
diffmask_model = diffmask_imagenet
image = torch.from_numpy(image).permute(2, 0, 1).float() / 255
dm_image = feature_extractor(image).unsqueeze(0)
dm_out = model.get_mask(dm_image)
mask = dm_out["mask"][0].detach()
pred = dm_out["pred_class"][0].detach()
pred = diffmask_model.model.config.id2label[pred.item()]
masked_img = draw_mask(image, mask)
heatmap = draw_heatmap(image, mask)
return np.hstack((masked_img, heatmap)), pred
# Launch demo interface
gr.Interface(
get_mask,
inputs=[gr.inputs.Image(label="Input", shape=(224, 224), source="upload", type="numpy"),
gr.inputs.Dropdown(["DiffMask-CiFAR-10", "DiffMask-ImageNet"])],
outputs=[gr.outputs.Image(label="Output"), gr.outputs.Label(label="Prediction")],
title="Vision DiffMask Demo",
live=True,
).launch()
|