ai-models/RMBG-2.0
2
0
Fork 0
mirror of https://www.modelscope.cn/AI-ModelScope/RMBG-2.0.git synced 2026-04-02 19:12:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update README.md

Browse Source
This commit is contained in:
Cherrytest
2025-09-22 07:12:51 +00:00
parent fdc5d45001
commit cf6aec8113
2 changed files with 4 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
README.md
View File

@ -146,15 +146,12 @@ transformers
```python ```python
from PIL import Image from PIL import Image
import matplotlib.pyplot as plt
import torch import torch
from torchvision import transforms from torchvision import transforms
from transformers import AutoModelForImageSegmentation from transformers import AutoModelForImageSegmentation
model = AutoModelForImageSegmentation.from_pretrained('briaai/RMBG-2.0', trust_remote_code=True) device = 'cuda' if torch.cuda.is_available() else 'cpu'
torch.set_float32_matmul_precision(['high', 'highest'][0]) model = AutoModelForImageSegmentation.from_pretrained('briaai/RMBG-2.0', trust_remote_code=True).eval().to(device)
model.to('cuda')
model.eval()
# Data settings # Data settings
image_size = (1024, 1024) image_size = (1024, 1024)
@ -165,7 +162,7 @@ transform_image = transforms.Compose([
]) ])
image = Image.open(input_image_path) image = Image.open(input_image_path)
input_images = transform_image(image).unsqueeze(0).to('cuda') input_images = transform_image(image).unsqueeze(0).to(device)
# Prediction # Prediction
with torch.no_grad(): with torch.no_grad():

1
birefnet.py
View File

@ -6,6 +6,7 @@ from transformers import PretrainedConfig
class Config(PretrainedConfig): class Config(PretrainedConfig):
def __init__(self) -> None: def __init__(self) -> None:
super().__init__()
# PATH settings # PATH settings
self.sys_home_dir = os.path.expanduser('~') # Make up your file system as: SYS_HOME_DIR/codes/dis/BiRefNet, SYS_HOME_DIR/datasets/dis/xx, SYS_HOME_DIR/weights/xx self.sys_home_dir = os.path.expanduser('~') # Make up your file system as: SYS_HOME_DIR/codes/dis/BiRefNet, SYS_HOME_DIR/datasets/dis/xx, SYS_HOME_DIR/weights/xx