Skip to content

Commit

Permalink
Added 'force_return_bytes' param to force bytes output. (#663)
Browse files Browse the repository at this point in the history
Fix 'Input type 'System.Bytes' is not supported'. When using .net, set  force_return_bytes to True to return bytes output.
  • Loading branch information
syaifulnizamyahya authored Aug 4, 2024
1 parent d1e0073 commit 90536b6
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 6 deletions.
15 changes: 15 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
14 changes: 8 additions & 6 deletions rembg/bg.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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)

Expand Down

0 comments on commit 90536b6

Please sign in to comment.