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 diff --git a/rembg/bg.py b/rembg/bg.py index 1f450179..fc99793b 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 diff --git a/rembg/cli.py b/rembg/cli.py index 983f7dd2..cd1ea84d 100644 --- a/rembg/cli.py +++ b/rembg/cli.py @@ -92,6 +92,14 @@ 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 +184,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(