zero-shot / app.py
quentin99's picture
Update app.py
bc12728
raw history blame
No virus
1.52 kB
from turtle import title
import gradio as gr
from transformers import pipeline
import numpy as np
from PIL import Image
pipes = {
"ViT/B-16": pipeline("zero-shot-image-classification", model="OFA-Sys/chinese-clip-vit-base-patch16")
}
inputs = [
gr.inputs.Image(type='pil',
label="Image 输入图片"),
gr.inputs.Textbox(lines=1,
label="Candidate Labels 候选分类标签"),
gr.inputs.Radio(choices=["ViT/B-16"], type="value", default="ViT/B-16", label="Model 模型规模"),
gr.inputs.Textbox(lines=1, label="Prompt Template Prompt模板 ({}指代候选标签)", default="一张{}的图片。"),
]
images="festival.jpg"
def shot(image, labels_text, model_name, hypothesis_template):
labels = [label.strip(" ") for label in labels_text.strip(" ").split(",")]
res = pipes[model_name](images=image,
candidate_labels=labels,
hypothesis_template=hypothesis_template)
return {dic["label"]: dic["score"] for dic in res}
iface = gr.Interface(shot,
inputs,
"label",
examples=[["festival.jpg", "灯笼, 鞭炮, 对联", "ViT/B-16", "一张{}的图片。"]],
description="""To play with this demo, add a picture and a list of labels in Chinese separated by commas. 上传图片,并输入多个分类标签,用英文逗号分隔。可点击页面最下方示例参考。""",
title="Zero-shot Image Classification (中文零样本图像分类)")
iface.launch()