浏览代码

Avoid overflow casting to uint8

PENGUINLIONG 4 月之前
父节点
当前提交
fdedf86e56
共有 1 个文件被更改,包括 1 次插入1 次删除
  1. 1 1
      rembg/sessions/u2net.py

+ 1 - 1
rembg/sessions/u2net.py

@@ -41,7 +41,7 @@ class U2netSession(BaseSession):
         pred = (pred - mi) / (ma - mi)
         pred = np.squeeze(pred)
 
-        mask = Image.fromarray((pred * 255).astype("uint8"), mode="L")
+        mask = Image.fromarray((pred.clip(0, 1) * 255).astype("uint8"), mode="L")
         mask = mask.resize(img.size, Image.Resampling.LANCZOS)
 
         return [mask]