app.py

app.py

@@ -0,0 +1,66 @@
+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
+
+# 加载Taiyi 中文 word encoder
+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的image encoder
+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是尺度系数
+        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 = np.around(probs, 3)[0]
+        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='自定义类别')
+                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)