|
|
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 !=""] |
|
|
|
|
|
|
|
|
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) |
|
|
|
|
|
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() |
