Merge pull request #414 from Flippchen/main

Add background color functionality to the remove function

rembg/bg.py

@@ -105,6 +105,14 @@ def post_process(mask: np.ndarray) -> np.ndarray:
     return mask
 
 
+def apply_background_color(img: PILImage, color: List[int]) -> PILImage:
+    r, g, b = color
+    colored_image = Image.new("RGBA", img.size, (r, g, b, 255))
+    colored_image.paste(img, mask=img)
+
+    return colored_image
+
+
 def remove(
     data: Union[bytes, PILImage, np.ndarray],
     alpha_matting: bool = False,
@@ -114,6 +122,7 @@ def remove(
     session: Optional[BaseSession] = None,
     only_mask: bool = False,
     post_process_mask: bool = False,
+    color: Optional[List[int]] = None,
 ) -> Union[bytes, PILImage, np.ndarray]:
     if isinstance(data, PILImage):
         return_type = ReturnType.PILLOW
@@ -161,6 +170,9 @@ def remove(
     if len(cutouts) > 0:
         cutout = get_concat_v_multi(cutouts)
 
+    if color is not None:
+        cutout = apply_background_color(cutout, color)
+
     if ReturnType.PILLOW == return_type:
         return cutout
 

rembg/cli.py

@@ -92,6 +92,14 @@ def main() -> None:
     show_default=True,
     help="post process the mask",
 )
[email protected](
+    "-c",
+    "--color",
+    default=None,
+    nargs=3,
+    type=int,
+    help="Background color (R G B) to replace the removed background with",
+)
 @click.argument(
     "input", default=(None if sys.stdin.isatty() else "-"), type=click.File("rb")
 )
@@ -176,6 +184,15 @@ def i(model: str, input: IO, output: IO, **kwargs) -> None:
     show_default=True,
     help="watches a folder for changes",
 )
[email protected](
+    "-c",
+    "--color",
+    default=None,
+    type=(int, int, int),
+    nargs=3,
+    metavar="R G B",
+    help="background color (RGB) to replace removed areas",
+)
 @click.argument(
     "input",
     type=click.Path(