OriLib commited on
Commit
2b99d49
1 Parent(s): e2ebf36

Update replace_bg/utilities.py

Browse files
Files changed (1) hide show
  1. replace_bg/utilities.py +2 -2
replace_bg/utilities.py CHANGED
@@ -37,10 +37,10 @@ def get_control_image_tensor(vae, image, mask)->torch.Tensor:
37
  control_tensor = torch.cat([control_latents, mask_resized], dim=1)
38
  return control_tensor
39
 
40
- def remove_bg_from_image(image_path: str)->Image.Image:
41
  from transformers import pipeline
42
  pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
43
- mask = pipe(image_path, return_mask = True) # outputs a pillow mask
44
  return mask
45
 
46
  def paste_fg_over_image(gen_image: Image.Image, orig_image: Image.Image, fg_mask: Image.Image)->Image.Image:
 
37
  control_tensor = torch.cat([control_latents, mask_resized], dim=1)
38
  return control_tensor
39
 
40
+ def remove_bg_from_image(image)->Image.Image:
41
  from transformers import pipeline
42
  pipe = pipeline("image-segmentation", model="briaai/RMBG-1.4", trust_remote_code=True)
43
+ mask = pipe(image, return_mask = True) # outputs a pillow mask
44
  return mask
45
 
46
  def paste_fg_over_image(gen_image: Image.Image, orig_image: Image.Image, fg_mask: Image.Image)->Image.Image: