From 6d5d0865fd031ca6c7d6a46bda063f8010d6952f Mon Sep 17 00:00:00 2001 From: Flippchen <91947480+Flippchen@users.noreply.github.com> Date: Wed, 29 Mar 2023 11:05:48 +0200 Subject: [PATCH 1/5] added image background feature --- rembg/bg.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/rembg/bg.py b/rembg/bg.py index 1f450179..a30bc678 100644 --- a/rembg/bg.py +++ b/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 From 6862af72d82432f5781bafc203e970fbe1b0e3c9 Mon Sep 17 00:00:00 2001 From: Flippchen <91947480+Flippchen@users.noreply.github.com> Date: Wed, 29 Mar 2023 11:09:04 +0200 Subject: [PATCH 2/5] fixed Lint --- rembg/bg.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/rembg/bg.py b/rembg/bg.py index a30bc678..fc99793b 100644 --- a/rembg/bg.py +++ b/rembg/bg.py @@ -107,7 +107,7 @@ def post_process(mask: np.ndarray) -> np.ndarray: 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 = Image.new("RGBA", img.size, (r, g, b, 255)) colored_image.paste(img, mask=img) return colored_image From 09726ab9192c843a86d2388615107c5e3fa8af86 Mon Sep 17 00:00:00 2001 From: Flippchen <91947480+Flippchen@users.noreply.github.com> Date: Wed, 29 Mar 2023 11:18:02 +0200 Subject: [PATCH 3/5] added source to readme --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index fb73d776..c89c276f 100644 --- a/README.md +++ b/README.md @@ -265,7 +265,7 @@ The available models are: - u2net_human_seg ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_human_seg.onnx), [source](https://github.com/xuebinqin/U-2-Net)): A pre-trained model for human segmentation. - u2net_cloth_seg ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_cloth_seg.onnx), [source](https://github.com/levindabhi/cloth-segmentation)): A pre-trained model for Cloths Parsing from human portrait. Here clothes are parsed into 3 category: Upper body, Lower body and Full body. - silueta ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/silueta.onnx), [source](https://github.com/xuebinqin/U-2-Net/issues/295)): Same as u2net but the size is reduced to 43Mb. -- isnet-general-use ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/isnet-general-use.onnx), [source](https://github.com/xuebinqin/U-2-Net/issues/295)): https://github.com/xuebinqin/DIS. +- isnet-general-use ([download](https://github.com/danielgatis/rembg/releases/download/v0.0.0/isnet-general-use.onnx), [source](https://github.com/xuebinqin/DIS)): A new pre-trained model for general use cases. ### Some differences between the models result From ab425e8f133cecfa3a03d5518df88557834f106b Mon Sep 17 00:00:00 2001 From: Flippchen <91947480+Flippchen@users.noreply.github.com> Date: Wed, 29 Mar 2023 13:20:30 +0200 Subject: [PATCH 4/5] added CLI Support --- rembg/cli.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/rembg/cli.py b/rembg/cli.py index 983f7dd2..96608f8e 100644 --- a/rembg/cli.py +++ b/rembg/cli.py @@ -92,6 +92,13 @@ def main() -> None: show_default=True, help="post process the mask", ) +@click.option( + "-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 +183,15 @@ def i(model: str, input: IO, output: IO, **kwargs) -> None: show_default=True, help="watches a folder for changes", ) +@click.option( + "-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( From f5837a17993b27d14fad3a3c28968c1b1fc9e55b Mon Sep 17 00:00:00 2001 From: Flippchen <91947480+Flippchen@users.noreply.github.com> Date: Wed, 29 Mar 2023 13:24:26 +0200 Subject: [PATCH 5/5] fix pylint --- rembg/cli.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/rembg/cli.py b/rembg/cli.py index 96608f8e..cd1ea84d 100644 --- a/rembg/cli.py +++ b/rembg/cli.py @@ -93,11 +93,12 @@ def main() -> None: help="post process the mask", ) @click.option( - "-c", "--color", + "-c", + "--color", default=None, nargs=3, type=int, - help="Background color (R G B) to replace the removed background with" + 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")