浏览代码

Merge pull request #276 from goecho/main

fix: Expected RGB image of shape (?, ?, 3), but image.shape is (?, ?, 4)
Daniel Gatis 3 年之前
父节点
当前提交
a11480edac
共有 1 个文件被更改,包括 4 次插入0 次删除
  1. 4 0
      rembg/bg.py

+ 4 - 0
rembg/bg.py

@@ -37,6 +37,10 @@ def alpha_matting_cutout(
     background_threshold: int,
     erode_structure_size: int,
 ) -> PILImage:
+
+    if img.mode == "RGBA" or img.mode == "CMYK":
+        img = img.convert("RGB")
+
     img = np.asarray(img)
     mask = np.asarray(mask)