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 11:02:56 +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
from PIL import Image
import matplotlib.pyplot as plt
import torch
from torchvision import transforms
from transformers import AutoModelForImageSegmentation
model = AutoModelForImageSegmentation.from_pretrained('briaai/RMBG-2.0', trust_remote_code=True)
torch.set_float32_matmul_precision(['high', 'highest'][0])
model.to('cuda')
model.eval()
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = AutoModelForImageSegmentation.from_pretrained('briaai/RMBG-2.0', trust_remote_code=True).eval().to(device)
# Data settings
image_size = (1024, 1024)
@ -165,7 +162,7 @@ transform_image = transforms.Compose([
])
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
with torch.no_grad():