briaai
/

RMBG-2.0

Image Segmentation

remove background

background-removal

legal liability

Model card Files Files and versions Community

OriLib commited on 11 days ago

Commit

1ea5224

•

1 Parent(s): 30a4470

Update README.md

Files changed (1) hide show

README.md +2 -2

README.md CHANGED Viewed

@@ -119,9 +119,9 @@ from PIL import Image
 import matplotlib.pyplot as plt
 import torch
 from torchvision import transforms
-from models.birefnet import BiRefNet
-birefnet = BiRefNet.from_pretrained('briaai/RMBG-2.0')
 torch.set_float32_matmul_precision(['high', 'highest'][0])
 birefnet.to('cuda')
 birefnet.eval()

 import matplotlib.pyplot as plt
 import torch
 from torchvision import transforms
+from transformers import AutoModelForImageSegmentation
+birefnet = AutoModelForImageSegmentation.from_pretrained('briaai/RMBG-2.0', trust_remote_code=True)
 torch.set_float32_matmul_precision(['high', 'highest'][0])
 birefnet.to('cuda')
 birefnet.eval()