integrate with huggingface (#3)

- integrate with huggingface (5ad69f0c63c1820bf9c9486fac0127ded199c980)


Co-authored-by: Hafedh Hichri <[email protected]>

.gitattributes

@@ -38,3 +38,4 @@ BEN2_demo_pictures/grid_example2.png filter=lfs diff=lfs merge=lfs -text
 BEN2_demo_pictures/grid_example3.png filter=lfs diff=lfs merge=lfs -text
 BEN2_demo_pictures/grid_example6.png filter=lfs diff=lfs merge=lfs -text
 BEN2_demo_pictures/grid_example7.png filter=lfs diff=lfs merge=lfs -text
+BEN2_demo_pictures/model_comparison.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,6 +1,7 @@
 ---
 license: mit
 pipeline_tag: image-segmentation
+library_name: ben2
 tags:
 - BEN2
 - background-remove
@@ -11,6 +12,8 @@ tags:
 - background
 - remove background
 - pytorch
+- model_hub_mixin
+- pytorch_model_hub_mixin
 ---
 
 # BEN2: Background Erase Network
@@ -37,10 +40,16 @@ BEN2 was trained on the DIS5k and our 22K proprietary segmentation dataset. Our
 - Follow us on X: https://x.com/PramaResearch/
 
 
-## Quick start code (inside cloned repo)
+## Installation
+
+```
+pip install -e "git+https://github.com/PramaLLC/BEN2.git#egg=ben2"
+```
+
+## Quick start code
 
 ```python
-import BEN2
+from ben2 import BEN_Base
 from PIL import Image
 import torch
 
@@ -49,9 +58,9 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 file = "./image.png" # input image
 
-model = BEN2.BEN_Base().to(device).eval() #init pipeline
+model = BEN_Base.from_pretrained("PramaLLC/BEN2")
+model.to(device).eval()
 
-model.loadcheckpoints("./BEN2_Base.pth")
 image = Image.open(file)
 foreground = model.inference(image, refine_foreground=False,) #Refine foreground is an extract postprocessing step that increases inference time but can improve matting edges. The default value is False.
 
@@ -63,7 +72,7 @@ foreground.save("./foreground.png")
 ## Batch image processing
 
 ```python
-import BEN2
+from ben2 import BEN_Base
 from PIL import Image
 import torch
 
@@ -72,9 +81,9 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 
 
-model = BEN2.BEN_Base().to(device).eval() #init pipeline
+model = BEN_Base.from_pretrained("PramaLLC/BEN2")
+model.to(device).eval()
 
-model.loadcheckpoints("./BEN2_Base.pth")
 
 file1 = "./image1.png" # input image1
 file2 = "./image2.png" # input image2
@@ -102,7 +111,7 @@ sudo apt install ffmpeg
 ```
 
 ```python
-import BEN2
+from ben2 import BEN_Base
 from PIL import Image
 import torch
 
@@ -111,11 +120,8 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 video_path = "/path_to_your_video.mp4"# input video
 
-model = BEN2.BEN_Base().to(device).eval() #init pipeline
-
-model.loadcheckpoints("./BEN2_Base.pth")
-
-
+model = BEN_Base.from_pretrained("PramaLLC/BEN2")
+model.to(device).eval()
 
 
 model.segment_video(
@@ -146,6 +152,3 @@ RMBG 2.0 did not preserve the DIS 5k validation dataset
 ![Example 7](BEN2_demo_pictures/grid_example7.png)
 
 
-## Installation
-1. Clone Repo
-2. Install requirements.txt

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ea8b7907176a09667c86343dc7d00de6a6d871076cb90bb5f753618fd6fb3ebb
+size 380577976