repository·dev·Indexed 43 days ago
https://github.com/bria-ai/rmbg-2.0A state-of-the-art image segmentation model that separates foreground objects from backgrounds by producing high-quality 8-bit grayscale alpha mattes. Available for non-commercial use under CC BY-NC 4.0 or via commercial license and API endpoints.
BRIA RMBG-2.0 is a dichotomous image segmentation model. Instead of a binary mask, the model outputs a single-channel 8-bit grayscale alpha matte.
Each pixel value in the output represents the opacity level of the corresponding pixel in the original image. This allows developers to apply custom thresholds for foreground-background separation, providing more control when integrating the model into complex image processing pipelines.
To run the model locally using the provided Python implementation, ensure you have the following dependencies installed:
torch
torchvision
pillow
kornia
transformersBRIA RMBG v2.0 is available through several distribution channels depending on your integration needs:
You can use the transformers library to load the briaai/RMBG-2.0 model. The workflow involves loading the model with trust_remote_code=True, preprocessing the image to a specific size (e.g., 1024x1024) and normalizing it, performing a forward pass, and then applying the resulting sigmoid alpha matte to the original image.
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()
# 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('cuda')
# 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")