let5see/IOPaint
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add remove_bg_device

Browse Source
This commit is contained in:
Qing
2024-11-23 15:51:05 +08:00
parent d29fe6ecb5
commit b7699a0f26
10 changed files with 64 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
iopaint/plugins/briarmbg.py
View File

@@ -480,7 +480,7 @@ def create_briarmbg_session():
return net
def briarmbg_process(bgr_np_image, session, only_mask=False):
def briarmbg_process(device, bgr_np_image, session, only_mask=False):
# prepare input
orig_bgr_image = Image.fromarray(bgr_np_image)
w, h = orig_im_size = orig_bgr_image.size
@@ -490,6 +490,7 @@ def briarmbg_process(bgr_np_image, session, only_mask=False):
im_tensor = torch.unsqueeze(im_tensor, 0)
im_tensor = torch.divide(im_tensor, 255.0)
im_tensor = normalize(im_tensor, [0.5, 0.5, 0.5], [1.0, 1.0, 1.0])
im_tensor = im_tensor.to(device)
# inference
result = session(im_tensor)
# post process