Spaces:
Sleeping
Sleeping
Commit
•
057eeaa
1
Parent(s):
7790ebf
added color
Browse files- example_inference.py +2 -2
example_inference.py
CHANGED
@@ -7,7 +7,7 @@ from huggingface_hub import hf_hub_download
|
|
7 |
import io as IO
|
8 |
import base64
|
9 |
|
10 |
-
def example_inference(im_path, transprent_bg=False):
|
11 |
|
12 |
net = BriaRMBG()
|
13 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
@@ -26,7 +26,7 @@ def example_inference(im_path, transprent_bg=False):
|
|
26 |
|
27 |
# post process
|
28 |
result_image = postprocess_image(result[0][0], orig_im_size)
|
29 |
-
bgColor = (0,0,0, 0) if transprent_bg else
|
30 |
# save result
|
31 |
pil_im = Image.fromarray(result_image)
|
32 |
no_bg_image = Image.new("RGBA", pil_im.size, bgColor)
|
|
|
7 |
import io as IO
|
8 |
import base64
|
9 |
|
10 |
+
def example_inference(im_path, transprent_bg=False, color=(255, 255, 255, 255)):
|
11 |
|
12 |
net = BriaRMBG()
|
13 |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
26 |
|
27 |
# post process
|
28 |
result_image = postprocess_image(result[0][0], orig_im_size)
|
29 |
+
bgColor = (0,0,0, 0) if transprent_bg else color
|
30 |
# save result
|
31 |
pil_im = Image.fromarray(result_image)
|
32 |
no_bg_image = Image.new("RGBA", pil_im.size, bgColor)
|