RMBG v1.4 是最先进的背景去除模型，旨在有效地将一系列类别和图像类型的前景与背景分开。该模型已在精心挑选的数据集上进行了训练，其中包括：一般库存图像、电子商务、游戏和广告内容，使其适用于支持大规模企业内容创建的商业用例。其准确性、效率和多功能性目前可与领先的 Source Available 型号相媲美。在内容安全、合法许可的数据集和偏见缓解至关重要的情况下，它是理想的选择。
模型链接：https://huggingface.co/briaai/RMBG-1.4
Demo 链接：https://huggingface.co/spaces/briaai/BRIA-RMBG-1.4
下面是测试的代码：
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
import streamlit as st from transformers import pipeline from PIL import Image import io from streamlit_image_comparison import image_comparison # 定义删除背景模型 def load_model(): return pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True) # 处理图像的函数 def process_image(pipe, image): # 处理图片：获取删除背景后的图像和 mask pillow_mask = pipe(image, return_mask=True) # 获取遮罩 pillow_image = pipe(image) # 应用遮罩后的图像 # 将结果保存为临时文件 mask_buffer = io.BytesIO() pillow_mask.save(mask_buffer, format="PNG") mask_buffer.seek(0) image_buffer = io.BytesIO() pillow_image.save(image_buffer, format="PNG") image_buffer.seek(0) return pillow_mask, pillow_image, mask_buffer, image_buffer # 主函数，用于运行 Streamlit 应用 def main(): # 加载模型 pipe = load_model() # Streamlit 页面设置 st.title("删除背景工具") st.write("请上传一张图片进行处理") # 上传图片文件 uploaded_file = st.file_uploader("选择一张图片", type=["jpg", "jpeg", "png"]) # 如果上传了图片，显示原图并提供删除背景按钮 if uploaded_file is not None: # 读取并显示原图 image = Image.open(uploaded_file) # 创建一列用于显示上传的文件和删除背景按钮 col1, _ = st.columns([2, 1]) with col1: st.subheader("原图") st.image(image, caption="上传的原图", use_column_width=True) # 将删除背景按钮放在图片下面 if st.button("删除背景"): # 调用删除背景的函数 pillow_mask, pillow_image, mask_buffer, image_buffer = process_image(pipe, image) # 将处理结果保存到 session_state 中 st.session_state["pillow_mask"] = pillow_mask st.session_state["pillow_image"] = pillow_image st.session_state["original_image"] = image st.session_state["mask_buffer"] = mask_buffer st.session_state["image_buffer"] = image_buffer # 如果已经处理过图像，显示删除背景前后对比 if "pillow_mask" in st.session_state and "pillow_image" in st.session_state: original_image = st.session_state["original_image"] processed_mask = st.session_state["pillow_mask"] processed_image = st.session_state["pillow_image"] # 创建两列用于显示删除背景后的 Mask 和 Output 图像 col1, col2 = st.columns(2) with col1: st.subheader("删除背景的 Mask") st.image(processed_mask, caption="删除背景的 Mask", use_column_width=True) st.download_button(label="下载遮罩图像", data=st.session_state["mask_buffer"], file_name="output_mask.png", mime="image/png") with col2: st.subheader("删除背景后的图像") st.image(processed_image, caption="删除背景后的图像", use_column_width=True) st.download_button(label="下载删除背景后的图像", data=st.session_state["image_buffer"], file_name="output_image.png", mime="image/png") # 提供下载链接 # 显示图像对比工具 st.subheader("删除背景前后对比") image_comparison( img1=original_image, img2=processed_image, label1="原图", label2="删除背景后图像" ) # 运行主函数 if __name__ == "__main__": main() |
实现简单的文件上传，会显示处理后的图像和 mask,以及对比情况，下面是样例图片
