RMBG-1.4 / README.md
Negev900's picture
Update README.md
8dc97d7
|
raw
history blame
2.68 kB
---
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](https://bria.ai/bria-2-0-huggingface-model-license-agreement/)
- **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](https://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](https://bria.ai/privacy-policy/) and [Terms & conditions](https://bria.ai/terms-and-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
```python
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)
```