OriLib commited on
Commit
3819229
1 Parent(s): aadede3

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +17 -15
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
- from glob import glob
99
- from tqdm import tqdm
100
- import cv2
101
  import torch.nn.functional as F
102
  from torchvision.transforms.functional import normalize
103
- from models import BriaRMBG
104
-
105
- input_size=[1024,1024]
106
- net=BriaRMBG()
107
 
108
  model_path = "./model.pth"
109
- im_path = "./example_image.jpg"
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=input_size, mode='bilinear').type(torch.uint8)
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
- cv2.imwrite(os.path.join(result_path, im_name+".png"), im_array)
 
 
 
 
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
  ```