File size: 3,117 Bytes
978d150 |
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 |
import gradio as gr
import torch
import io
from PIL import Image
from transformers import (
AutoImageProcessor,
AutoTokenizer,
AutoModelForCausalLM,
)
import numpy as np
model_root = "qihoo360/fg-clip-base"
model = AutoModelForCausalLM.from_pretrained(model_root,trust_remote_code=True)
device = model.device
tokenizer = AutoTokenizer.from_pretrained(model_root)
image_processor = AutoImageProcessor.from_pretrained(model_root)
import math
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
def Get_Densefeature(image, candidate_labels):
"""
Takes an image and a comma-separated string of candidate labels,
and returns the classification scores.
"""
candidate_labels = [label.lstrip(" ") for label in candidate_labels.split(",") if label !=""]
# print(candidate_labels)
image_size=224
image = image.convert("RGB")
image = image.resize((image_size,image_size))
image_input = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].to(device)
with torch.no_grad():
dense_image_feature = model.get_image_dense_features(image_input)
captions = [candidate_labels[0]]
caption_input = torch.tensor(tokenizer(captions, max_length=77, padding="max_length", truncation=True).input_ids, dtype=torch.long, device=device)
text_feature = model.get_text_features(caption_input,walk_short_pos=True)
text_feature = text_feature / text_feature.norm(p=2, dim=-1, keepdim=True)
dense_image_feature = dense_image_feature / dense_image_feature.norm(p=2, dim=-1, keepdim=True)
similarity = dense_image_feature.squeeze() @ text_feature.squeeze().T
similarity = similarity.cpu().numpy()
patch_size = int(math.sqrt(similarity.shape[0]))
original_shape = (patch_size, patch_size)
show_image = similarity.reshape(original_shape)
fig = plt.figure(figsize=(6, 6))
plt.imshow(show_image)
plt.title('similarity Visualization')
plt.axis('off')
buf = io.BytesIO()
plt.savefig(buf, format='png')
buf.seek(0)
plt.close(fig)
pil_img = Image.open(buf)
# buf.close()
return pil_img
with gr.Blocks() as demo:
gr.Markdown("# FG-CLIP Densefeature")
gr.Markdown(
"This app uses the FG-CLIP model (qihoo360/fg-clip-base) for Densefeature show on CPU :"
)
gr.Markdown(
"<span style='color: red; font-weight: bold;'>⚠️ (Run DenseFeature) only support one class</span>"
)
with gr.Row():
with gr.Column():
image_input = gr.Image(type="pil")
text_input = gr.Textbox(label="Input a label")
dfs_button = gr.Button("Run Densefeature", visible=True)
with gr.Column():
dfs_output = gr.Image(label="Similarity Visualization", type="pil")
examples = [
["./cat_dfclor.jpg", "white cat,"],
]
gr.Examples(
examples=examples,
inputs=[image_input, text_input],
)
dfs_button.click(fn=Get_Densefeature, inputs=[image_input, text_input], outputs=dfs_output)
demo.launch() |