Ver Fonte

Avoid overflow casting to uint8

PENGUINLIONG há 4 meses atrás
pai
commit
fdedf86e56
1 ficheiros alterados com 1 adições e 1 exclusões
  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]