|
from PIL import Image |
|
import torch |
|
from transformers import BertForSequenceClassification, BertConfig, BertTokenizer |
|
from transformers import CLIPProcessor, CLIPModel |
|
import numpy as np |
|
import time |
|
import gradio as gr |
|
import re |
|
|
|
|
|
text_tokenizer = BertTokenizer.from_pretrained("IDEA-CCNL/Taiyi-CLIP-Roberta-102M-Chinese") |
|
text_encoder = BertForSequenceClassification.from_pretrained("IDEA-CCNL/Taiyi-CLIP-Roberta-102M-Chinese").eval() |
|
|
|
clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") |
|
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32") |
|
|
|
|
|
def imgclassfiy(query_texts,img_url): |
|
start_time = time.time() |
|
query_texts =re.split(",|,",query_texts) |
|
text = text_tokenizer(query_texts, return_tensors='pt', padding=True)['input_ids'] |
|
url = img_url |
|
|
|
image = processor(images=Image.open(url), return_tensors="pt") |
|
|
|
with torch.no_grad(): |
|
image_features = clip_model.get_image_features(**image) |
|
text_features = text_encoder(text).logits |
|
|
|
|
|
image_features = image_features / image_features.norm(dim=1, keepdim=True) |
|
text_features = text_features / text_features.norm(dim=1, keepdim=True) |
|
|
|
|
|
logit_scale = clip_model.logit_scale.exp() |
|
logits_per_image = logit_scale * image_features @ text_features.t() |
|
logits_per_text = logits_per_image.t() |
|
probs = logits_per_image.softmax(dim=-1).cpu().numpy() |
|
|
|
|
|
res = query_texts[np.argmax(probs)] |
|
|
|
end_time = time.time() |
|
print('用时:', end_time - start_time) |
|
return res |
|
|
|
if __name__ =="__main__": |
|
|
|
with gr.Blocks(title="自定义类别的图像分类") as demo: |
|
|
|
gr.HTML('<br>') |
|
gr.HTML( |
|
f'<center><p style="color:#4377ec;font-size:42px;font-weight:bold;text-shadow: #FDEDB7 2px 0 0, #FDEDB7 0 2px 0, #FDEDB7 -2px 0 0, #FDEDB7 0 -2px 0;">自定义类别的图像分类</p></center>') |
|
gr.HTML('<br>') |
|
with gr.Row() as row: |
|
with gr.Column(): |
|
img_input = gr.Image(type="filepath") |
|
out_input = gr.Textbox(lable='自定义类别',placeholder='输入自定义类别,例如:猫,狗,兔子') |
|
text_btn = gr.Button("提交") |
|
|
|
with gr.Column(scale=5): |
|
img_out = gr.Textbox(lable='输出类别') |
|
|
|
text_btn.click(fn=imgclassfiy, inputs=[out_input,img_input], outputs=[img_out]) |
|
|
|
demo.launch(show_api=False,inbrowser=True) |