diff --git a/README.md b/README.md index a9d1472a..ea836f06 100644 --- a/README.md +++ b/README.md @@ -274,6 +274,21 @@ output = remove(input) cv2.imwrite(output_path, output) ``` +Force output as bytes + +```python +from rembg import remove + +input_path = 'input.png' +output_path = 'output.png' + +with open(input_path, 'rb') as i: + with open(output_path, 'wb') as o: + input = i.read() + output = remove(input, force_return_bytes=True) + o.write(output) +``` + How to iterate over files in a performatic way ```python diff --git a/rembg/bg.py b/rembg/bg.py index c0ccd783..51f2ff9d 100644 --- a/rembg/bg.py +++ b/rembg/bg.py @@ -213,13 +213,14 @@ def remove( only_mask: bool = False, post_process_mask: bool = False, bgcolor: Optional[Tuple[int, int, int, int]] = None, + force_return_bytes: bool = False, *args: Optional[Any], **kwargs: Optional[Any] ) -> Union[bytes, PILImage, np.ndarray]: """ Remove the background from an input image. - This function takes in various parameters and returns a modified version of the input image with the background removed. The function can handle input data in the form of bytes, a PIL image, or a numpy array. The function first checks the type of the input data and converts it to a PIL image if necessary. It then fixes the orientation of the image and proceeds to perform background removal using the 'u2net' model. The result is a list of binary masks representing the foreground objects in the image. These masks are post-processed and combined to create a final cutout image. If a background color is provided, it is applied to the cutout image. The function returns the resulting cutout image in the format specified by the input 'return_type' parameter. + This function takes in various parameters and returns a modified version of the input image with the background removed. The function can handle input data in the form of bytes, a PIL image, or a numpy array. The function first checks the type of the input data and converts it to a PIL image if necessary. It then fixes the orientation of the image and proceeds to perform background removal using the 'u2net' model. The result is a list of binary masks representing the foreground objects in the image. These masks are post-processed and combined to create a final cutout image. If a background color is provided, it is applied to the cutout image. The function returns the resulting cutout image in the format specified by the input 'return_type' parameter or as python bytes if force_return_bytes is true. Parameters: data (Union[bytes, PILImage, np.ndarray]): The input image data. @@ -231,23 +232,24 @@ def remove( only_mask (bool, optional): Flag indicating whether to return only the binary masks. Defaults to False. post_process_mask (bool, optional): Flag indicating whether to post-process the masks. Defaults to False. bgcolor (Optional[Tuple[int, int, int, int]], optional): Background color for the cutout image. Defaults to None. + force_return_bytes (bool, optional): Flag indicating whether to return the cutout image as bytes. Defaults to False. *args (Optional[Any]): Additional positional arguments. **kwargs (Optional[Any]): Additional keyword arguments. Returns: Union[bytes, PILImage, np.ndarray]: The cutout image with the background removed. """ - if isinstance(data, PILImage): - return_type = ReturnType.PILLOW - img = data - elif isinstance(data, bytes): + if isinstance(data, bytes) or force_return_bytes: return_type = ReturnType.BYTES img = Image.open(io.BytesIO(data)) + elif isinstance(data, PILImage): + return_type = ReturnType.PILLOW + img = data elif isinstance(data, np.ndarray): return_type = ReturnType.NDARRAY img = Image.fromarray(data) else: - raise ValueError("Input type {} is not supported.".format(type(data))) + raise ValueError("Input type {} is not supported. Try using force_return_bytes=True to force python bytes output".format(type(data))) putalpha = kwargs.pop("putalpha", False)