Update weights.

BiRefNet-massive-epoch_260.pth → BiRefNet-portrait-TR_P3M_10k-epoch_120.pth

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:92bb252307a56c9bbf95a258a96f60e96a8ccbb8d45275de0b8c76dbda0cd42a
+oid sha256:9a31ed736e17bc36a7ab539b71419b27f69b839e8ba132c6a67fe77edceb4141
 size 885038517

app.py

@@ -41,7 +41,7 @@ class ImagePreprocessor():
 
 
 model = BiRefNet(bb_pretrained=False)
-state_dict = ['./BiRefNet-DIS-epoch_590.pth', 'BiRefNet-massive-epoch_240.pth', 'BiRefNet-massive-epoch_260.pth'][-1]
+state_dict = ['./BiRefNet-DIS-epoch_590.pth', 'BiRefNet-massive-epoch_240.pth', 'BiRefNet-portrait-TR_P3M_10k-epoch_120.pth'][1]
 if os.path.exists(state_dict):
     birefnet_dict = torch.load(state_dict, map_location="cpu")
     unwanted_prefix = '_orig_mod.'
@@ -81,7 +81,7 @@ def predict(image, resolution):
     for image, pred in zip(images, preds):
         image = image.resize(pred.shape[::-1])
         pred = np.repeat(np.expand_dims(pred, axis=-1), 3, axis=-1)
-        image_preds.append((pred * image).astype(np.uint8))
+        image_preds.append(((pred>0.5) * image).astype(np.uint8))
 
     return image, image_preds[0]