Update README.md
Browse files
README.md
CHANGED
@@ -88,41 +88,40 @@ For clarity, we provide our data distribution according to different categories,
|
|
88 |
|
89 |
RMBG v1.4 is developed on the [IS-Net](https://github.com/xuebinqin/DIS) enhanced with our unique training scheme and proprietary dataset. These modifications significantly improve the model’s accuracy and effectiveness in diverse image-processing scenarios.
|
90 |
|
|
|
|
|
|
|
|
|
91 |
|
92 |
## Usage
|
93 |
|
94 |
```python
|
95 |
-
import os
|
96 |
import numpy as np
|
97 |
from skimage import io
|
98 |
-
|
99 |
-
from tqdm import tqdm
|
100 |
-
import cv2
|
101 |
import torch.nn.functional as F
|
102 |
from torchvision.transforms.functional import normalize
|
103 |
-
from
|
104 |
-
|
105 |
-
input_size=[1024,1024]
|
106 |
-
net=BriaRMBG()
|
107 |
|
108 |
model_path = "./model.pth"
|
109 |
-
im_path = "./
|
110 |
-
result_path = "."
|
111 |
|
|
|
112 |
if torch.cuda.is_available():
|
113 |
-
net.load_state_dict(torch.load(model_path))
|
114 |
-
net=net.cuda()
|
115 |
else:
|
116 |
net.load_state_dict(torch.load(model_path,map_location="cpu"))
|
117 |
net.eval()
|
118 |
|
119 |
# prepare input
|
|
|
120 |
im = io.imread(im_path)
|
121 |
if len(im.shape) < 3:
|
122 |
im = im[:, :, np.newaxis]
|
123 |
im_size=im.shape[0:2]
|
124 |
im_tensor = torch.tensor(im, dtype=torch.float32).permute(2,0,1)
|
125 |
-
im_tensor = F.interpolate(torch.unsqueeze(im_tensor,0), size=
|
126 |
image = torch.divide(im_tensor,255.0)
|
127 |
image = normalize(image,[0.5,0.5,0.5],[1.0,1.0,1.0])
|
128 |
|
@@ -139,7 +138,10 @@ mi = torch.min(result)
|
|
139 |
result = (result-mi)/(ma-mi)
|
140 |
|
141 |
# save result
|
142 |
-
im_name=im_path.split('/')[-1].split('.')[0]
|
143 |
im_array = (result*255).permute(1,2,0).cpu().data.numpy().astype(np.uint8)
|
144 |
-
|
|
|
|
|
|
|
|
|
145 |
```
|
|
|
88 |
|
89 |
RMBG v1.4 is developed on the [IS-Net](https://github.com/xuebinqin/DIS) enhanced with our unique training scheme and proprietary dataset. These modifications significantly improve the model’s accuracy and effectiveness in diverse image-processing scenarios.
|
90 |
|
91 |
+
## installation
|
92 |
+
git clone https://huggingface.co/briaai/RMBG-1.4
|
93 |
+
cd RMBG-1.4/
|
94 |
+
pip install -r requirements.txt
|
95 |
|
96 |
## Usage
|
97 |
|
98 |
```python
|
|
|
99 |
import numpy as np
|
100 |
from skimage import io
|
101 |
+
import torch
|
|
|
|
|
102 |
import torch.nn.functional as F
|
103 |
from torchvision.transforms.functional import normalize
|
104 |
+
from briarmbg import BriaRMBG
|
105 |
+
from PIL import Image
|
|
|
|
|
106 |
|
107 |
model_path = "./model.pth"
|
108 |
+
im_path = "./example_input.jpg"
|
|
|
109 |
|
110 |
+
net = BriaRMBG()
|
111 |
if torch.cuda.is_available():
|
112 |
+
net.load_state_dict(torch.load(model_path)).cuda()
|
|
|
113 |
else:
|
114 |
net.load_state_dict(torch.load(model_path,map_location="cpu"))
|
115 |
net.eval()
|
116 |
|
117 |
# prepare input
|
118 |
+
model_input_size=[1024,1024]
|
119 |
im = io.imread(im_path)
|
120 |
if len(im.shape) < 3:
|
121 |
im = im[:, :, np.newaxis]
|
122 |
im_size=im.shape[0:2]
|
123 |
im_tensor = torch.tensor(im, dtype=torch.float32).permute(2,0,1)
|
124 |
+
im_tensor = F.interpolate(torch.unsqueeze(im_tensor,0), size=model_input_size, mode='bilinear').type(torch.uint8)
|
125 |
image = torch.divide(im_tensor,255.0)
|
126 |
image = normalize(image,[0.5,0.5,0.5],[1.0,1.0,1.0])
|
127 |
|
|
|
138 |
result = (result-mi)/(ma-mi)
|
139 |
|
140 |
# save result
|
|
|
141 |
im_array = (result*255).permute(1,2,0).cpu().data.numpy().astype(np.uint8)
|
142 |
+
pil_im = Image.fromarray(np.squeeze(im_array))
|
143 |
+
no_bg_image = Image.new("RGBA", pil_im.size, (0,0,0,0))
|
144 |
+
orig_image = Image.open(im_path)
|
145 |
+
no_bg_image.paste(orig_image, mask=pil_im)
|
146 |
+
no_bg_image.save("example_image_no_bg.png")
|
147 |
```
|