2
0
Эх сурвалжийг харах

chore: Refactor remove function to handle unsupported input types

Daniel Gatis 1 жил өмнө
parent
commit
688b34841b
1 өөрчлөгдсөн 6 нэмэгдсэн , 2 устгасан
  1. 6 2
      rembg/bg.py

+ 6 - 2
rembg/bg.py

@@ -241,7 +241,7 @@ def remove(
     """
     if isinstance(data, bytes) or force_return_bytes:
         return_type = ReturnType.BYTES
-        img = Image.open(io.BytesIO(data))
+        img = Image.open(io.BytesIO(cast(bytes, data)))
     elif isinstance(data, PILImage):
         return_type = ReturnType.PILLOW
         img = data
@@ -249,7 +249,11 @@ def remove(
         return_type = ReturnType.NDARRAY
         img = Image.fromarray(data)
     else:
-        raise ValueError("Input type {} is not supported. Try using force_return_bytes=True to force python bytes output".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)