metadata

license: other
tags:
  - background-removal
  - Pytorch
  - vision

BRIA Background Removal v1.4 Model Card

100% automatically Background removal capability across all categories and image types that capture the variety of the world.

Model Description

Developed by: BRIA AI
Model type: Background removal image-to-image model
License: bria-2.0
Model Description: BRIA RMBG 1.4 is an image-to-image model trained exclusively on a professional-grade. It is designed for commercial.
Resources for more information: BRIA AI

Get Access

BRIA RMBG 1.3 is available under the BRIA RMBG 1.4 License Agreement, allowing commercial usage with an attribution model that supports our data contributors. To access the model, please contact us. By submitting this form, you agree to BRIA’s Privacy policy and Terms & conditions.

Training data

Bria-RMBG model was trained over 12000 high quality, high resolution images. All images were manualy labeled pixel-wise accuratly. The images belong to veriety of categories, the majority of them inclues people.

Qualitative Evaluation

Usage

import os
import numpy as np
from skimage import io
from glob import glob
from tqdm import tqdm
import cv2
import torch.nn.functional as F
from torchvision.transforms.functional import normalize
from models import BriaRMBG

input_size=[1024,1024]
net=BriaRMBG()

model_path = "./model.pth"
im_path = "./example_image.jpg"
result_path = "."

if torch.cuda.is_available():
    net.load_state_dict(torch.load(model_path))
    net=net.cuda()
else:
    net.load_state_dict(torch.load(model_path,map_location="cpu"))
net.eval()    

# prepare input
im = io.imread(im_path)
if len(im.shape) < 3:
    im = im[:, :, np.newaxis]
im_size=im.shape[0:2]
im_tensor = torch.tensor(im, dtype=torch.float32).permute(2,0,1)
im_tensor = F.interpolate(torch.unsqueeze(im_tensor,0), size=input_size, mode='bilinear').type(torch.uint8)
image = torch.divide(im_tensor,255.0)
image = normalize(image,[0.5,0.5,0.5],[1.0,1.0,1.0])

if torch.cuda.is_available():
    image=image.cuda()

# inference 
result=net(image)

# post process
result = torch.squeeze(F.interpolate(result[0][0], size=im_size, mode='bilinear') ,0)
ma = torch.max(result)
mi = torch.min(result)
result = (result-mi)/(ma-mi)

# save result
im_name=im_path.split('/')[-1].split('.')[0]
im_array = (result*255).permute(1,2,0).cpu().data.numpy().astype(np.uint8)
cv2.imwrite(os.path.join(result_path, im_name+".png"), im_array)