chore: Refactor remove function to handle unsupported input types

danielgatis · Aug 4, 2024 · 688b348 · 688b348
1 parent 90536b6
commit 688b348
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/rembg/bg.py b/rembg/bg.py
@@ -241,15 +241,19 @@ 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
     elif isinstance(data, np.ndarray):
         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)