update

app.py

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