init

app-gradio.py

@@ -0,0 +1,63 @@
+import gradio as gr
+import torch
+import io
+import base64
+import urllib.request
+from PIL import Image
+from process import process
+
+# 设备检测
+DEVICE = "GPU" if torch.cuda.is_available() else "CPU"
+
+
+def load_image(image, url):
+    """加载用户上传或URL图片"""
+    if image is not None:
+        return image
+    elif url:
+        try:
+            if url.startswith("http"):
+                with urllib.request.urlopen(url) as response:
+                    image_data = response.read()
+                    return Image.open(io.BytesIO(image_data))
+            elif url.startswith("data:image/"):
+                header, base64_data = url.split(",", 1)
+                return Image.open(io.BytesIO(base64.b64decode(base64_data)))
+        except Exception as e:
+            return None
+    return None
+
+
+def remove_background(image):
+    """移除背景"""
+    if image is None:
+        return None, None
+    mask, image_nbg = process(image)
+    return mask, image_nbg
+
+
+def interface(image, url):
+    """完整的Gradio处理流程"""
+    image = load_image(image, url)
+    if image is None:
+        return None, None, "请上传有效图片或输入正确的URL"
+    mask, image_nbg = remove_background(image)
+    return mask, image_nbg, "处理完成" if mask else "处理失败"
+
+
+# Gradio UI
+demo = gr.Interface(
+    fn=interface,
+    inputs=[gr.Image(type="pil", label="上传图片"), gr.Textbox(label="或输入图片URL")],
+    outputs=[
+        gr.Image(type="pil", label="掩码"),
+        gr.Image(type="pil", label="去除背景的图片"),
+    ],
+    title="AI 抠图 (RMBG 2.0)",
+    description="上传图片或提供URL，自动去除背景",
+    theme="default",
+    flagging_mode="never",
+)
+
+demo.queue()
+demo.launch()

process.py

@@ -0,0 +1,48 @@
+import os
+import torch
+import streamlit as st
+from PIL import Image
+from torchvision import transforms
+from transformers import AutoModelForImageSegmentation
+
+
+@st.cache_resource
+def load_model(model_id_or_path="briaai/RMBG-2.0", precision=0, device="cuda"):
+    model = AutoModelForImageSegmentation.from_pretrained(
+        model_id_or_path, trust_remote_code=True
+    )
+    torch.set_float32_matmul_precision(["high", "highest"][precision])
+    model.to(device)
+    _ = model.eval()
+
+    # Data settings
+    image_size = (1024, 1024)
+    transform_image = transforms.Compose(
+        [
+            transforms.Resize(image_size),
+            transforms.ToTensor(),
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+        ]
+    )
+
+    return model, transform_image
+
+
+def process(image: Image.Image) -> Image.Image:
+    if "RMBG-2.0" not in os.listdir("."):
+        os.system(
+            "modelscope download --model AI-ModelScope/RMBG-2.0 --local_dir RMBG-2.0 --exclude *.onnx *.bin"
+        )
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    precision = 0
+    model, transform = load_model("RMBG-2.0", precision=precision, device=device)
+    image = image.copy()
+    input_images = transform(image).unsqueeze(0).to(device)
+    with torch.no_grad():
+        preds = model(input_images)[-1].sigmoid().cpu()
+        pred = preds[0].squeeze()
+    pred_pil = transforms.ToPILImage()(pred)
+    mask = pred_pil.resize(image.size)
+    image.putalpha(mask)
+
+    return mask, image

requirements.txt

@@ -0,0 +1,11 @@
+torch
+torchvision
+pillow
+kornia
+transformers
+streamlit
+huggingface
+timm
+modelscope
+psutil
+gradio