目录RMBG-2.0模型下载推理示例代码rembg安装RMBG-2.0模型下载pip install atomgit atomgit download hf_mirrors/briaai/RMBG-2.0 -d ./推理示例代码from PIL import Image import torch from torchvision import transforms from transformers import AutoModelForImageSegmentation device cuda if torch.cuda.is_available() else cpu model AutoModelForImageSegmentation.from_pretrained(briaai/RMBG-2.0, trust_remote_codeTrue).eval().to(device) # 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]) ]) image Image.open(input_image_path) input_images transform_image(image).unsqueeze(0).to(device) # Prediction 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) image.save(no_bg_image.png)rembg安装pip install rembg[gpu]pip install rembgimport os from glob import glob from rembg import remove from argparse import ArgumentParser from PIL import Image if __name__ __main__: parser ArgumentParser() parser.add_argument(--path, typestr, requiredTrue, helpPath to input images) args parser.parse_args() imgs glob(os.path.join(args.path, *.png)) glob(os.path.join(args.path, *.jpg)) for img in imgs: path os.path.dirname(img) name os.path.basename(img).split(.)[0] _rmbg.png img_np Image.open(img) img remove(img_np) img.save(os.path.join(args.path, name))