added color

example_inference.py

@@ -7,7 +7,7 @@ from huggingface_hub import hf_hub_download
 import io as IO
 import base64
 
-def example_inference(im_path, transprent_bg=False):
+def example_inference(im_path, transprent_bg=False, color=(255, 255, 255, 255)):
 
     net = BriaRMBG()
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -26,7 +26,7 @@ def example_inference(im_path, transprent_bg=False):
 
     # post process
     result_image = postprocess_image(result[0][0], orig_im_size)
-    bgColor = (0,0,0, 0) if transprent_bg else (255,255,255, 255)
+    bgColor = (0,0,0, 0) if transprent_bg else color
     # save result
     pil_im = Image.fromarray(result_image)
     no_bg_image = Image.new("RGBA", pil_im.size, bgColor)