Spaces:
Running
Running
File size: 1,570 Bytes
696f3f1 |
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 |
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
|