فهرست منبع

Added 'force_return_bytes' param to force bytes output. (#663)

Fix 'Input type 'System.Bytes' is not supported'. When using .net, set  force_return_bytes to True to return bytes output.
syaifulnizamyahya 1 سال پیش
والد
کامیت
90536b6858
2فایلهای تغییر یافته به همراه23 افزوده شده و 6 حذف شده
  1. 15 0
      README.md
  2. 8 6
      rembg/bg.py

+ 15 - 0
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

+ 8 - 6
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)