Spaces:
Sleeping
Sleeping
import gradio as gr | |
import rembg | |
import functools | |
from gradio.mix import Parallel | |
from transformers import pipeline | |
from sentence_transformers import SentenceTransformer, util | |
generator = pipeline('text-generation', model='gpt2') | |
model_txt_embedding = SentenceTransformer("sentence-transformers/all-mpnet-base-v2") | |
# 文本向量 | |
def embedding(txt): | |
return str(model_txt_embedding.encode(txt, convert_to_tensor=True, convert_to_numpy=True).tolist()) | |
# 文本相似度 | |
def embedding_sim(text1, text2): | |
embeddings1 = model_txt_embedding.encode(text1, convert_to_tensor=True) | |
embeddings2 = model_txt_embedding.encode(text2, convert_to_tensor=True) | |
return str(util.cos_sim(embeddings1, embeddings2).item()) | |
# 抠图 | |
remove_bg_models = { | |
"常规": "u2net", | |
"人物": "u2net_human_seg", | |
"衣服": "u2net_cloth_seg" | |
} | |
def get_remove_session(model): | |
return rembg.new_session(model) | |
def remove_bg(image, model_cn, only_mask): | |
# return image | |
session = get_remove_session(remove_bg_models.get(model_cn)) | |
return rembg.remove(image, session=session, only_mask=only_mask) | |
tab1 = gr.Interface(fn=embedding, inputs="text", outputs="text") | |
tab2 = gr.Interface(fn=embedding_sim, inputs=["text", "text"], outputs="text") | |
tab3 = gr.Interface(fn=remove_bg, inputs=[ | |
"image", | |
gr.Dropdown(label="模型", choices=list(remove_bg_models.keys()), value="常规"), | |
gr.Checkbox(label="只返回蒙版", value=False) | |
], outputs="image") | |
app = gr.TabbedInterface([tab1, tab2, tab3], ["文本向量", "文本相似度", "抠图"]) | |
app.queue() | |
app.launch() |