README.md
5.0 KB · 154 lines · markdown Raw
1 ---
2 license: mit
3 pipeline_tag: image-segmentation
4 library_name: ben2
5 tags:
6 - BEN2
7 - background-remove
8 - mask-generation
9 - Dichotomous image segmentation
10 - background remove
11 - foreground
12 - background
13 - remove background
14 - pytorch
15 - model_hub_mixin
16 - pytorch_model_hub_mixin
17 - background removal
18 - background-removal
19 ---
20
21 # BEN2: Background Erase Network
22
23 [![arXiv](https://img.shields.io/badge/arXiv-2501.06230-b31b1b.svg)](https://arxiv.org/abs/2501.06230)
24 [![GitHub](https://img.shields.io/badge/GitHub-BEN2-black.svg)](https://github.com/PramaLLC/BEN2/)
25 [![Website](https://img.shields.io/badge/Website-backgrounderase.net-104233)](https://backgrounderase.net)
26
27 ## Overview
28 BEN2 (Background Erase Network) introduces a novel approach to foreground segmentation through its innovative Confidence Guided Matting (CGM) pipeline. The architecture employs a refiner network that targets and processes pixels where the base model exhibits lower confidence levels, resulting in more precise and reliable matting results. This model is built on BEN:
29 [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/ben-using-confidence-guided-matting-for/dichotomous-image-segmentation-on-dis-vd)](https://paperswithcode.com/sota/dichotomous-image-segmentation-on-dis-vd?p=ben-using-confidence-guided-matting-for)
30
31
32
33
34 ## BEN2 access
35 BEN2 was trained on the DIS5k and our 22K proprietary segmentation dataset. Our enhanced model delivers superior performance in hair matting, 4K processing, object segmentation, and edge refinement. Our Base model is open source. To try the full model through our free web demo or integrate BEN2 into your project with our API:
36 - 🌐 [backgrounderase.com](https://backgrounderase.com)
37
38
39 ## Contact us
40 - For access to our commercial model email us at sales@backgrounderase.com
41 - Our website: https://backgrounderase.com/
42 - Follow us on X: https://x.com/PramaResearch/
43
44
45 ## Installation
46
47 ```
48 pip install -e "git+https://github.com/PramaLLC/BEN2.git#egg=ben2"
49 ```
50
51 ## Quick start code
52
53 ```python
54 from ben2 import BEN_Base
55 from PIL import Image
56 import torch
57
58
59 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
60
61 file = "./image.png" # input image
62
63 model = BEN_Base.from_pretrained("PramaLLC/BEN2")
64 model.to(device).eval()
65
66 image = Image.open(file)
67 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.
68
69 foreground.save("./foreground.png")
70
71 ```
72
73
74 ## Batch image processing
75
76 ```python
77 from ben2 import BEN_Base
78 from PIL import Image
79 import torch
80
81
82 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
83
84
85
86 model = BEN_Base.from_pretrained("PramaLLC/BEN2")
87 model.to(device).eval()
88
89
90 file1 = "./image1.png" # input image1
91 file2 = "./image2.png" # input image2
92 image1 = Image.open(file1)
93 image2 = Image.open(file2)
94
95
96
97 foregrounds = model.inference([image1, image2]) # We recommend that the batch size not exceed 3 for consumer GPUs as there are minimal inference gains due to our custom batch processing for the MVANet decoding steps.
98 foregrounds[0].save("./foreground1.png")
99 foregrounds[1].save("./foreground2.png")
100
101 ```
102
103
104
105 # BEN2 video segmentation
106 [![BEN2 Demo](https://img.youtube.com/vi/skEXiIHQcys/0.jpg)](https://www.youtube.com/watch?v=skEXiIHQcys)
107
108 ## Video Segmentation
109
110 ```bash
111 sudo apt update
112 sudo apt install ffmpeg
113 ```
114
115 ```python
116 from ben2 import BEN_Base
117 from PIL import Image
118 import torch
119
120
121 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
122
123 video_path = "/path_to_your_video.mp4"# input video
124
125 model = BEN_Base.from_pretrained("PramaLLC/BEN2")
126 model.to(device).eval()
127
128
129 model.segment_video(
130 video_path= video_path,
131 output_path="./", # Outputs will be saved as foreground.webm or foreground.mp4. The default value is "./"
132 fps=0, # If this is set to 0 CV2 will detect the fps in the original video. The default value is 0.
133 refine_foreground=False, #refine foreground is an extract postprocessing step that increases inference time but can improve matting edges. The default value is False.
134 batch=1, # We recommended that batch size not exceed 3 for consumer GPUs as there are minimal inference gains. The default value is 1.
135 print_frames_processed=True, #Informs you what frame is being processed. The default value is True.
136 webm = False, # This will output an alpha layer video but this defaults to mp4 when webm is false. The default value is False.
137 rgb_value= (0, 255, 0) # If you do not use webm this will be the RGB value of the resulting background only when webm is False. The default value is a green background (0,255,0).
138 )
139
140
141 ```
142
143
144
145 **# BEN2 evaluation**
146 ![Model Comparison](BEN2_demo_pictures/model_comparison.png)
147
148 RMBG 2.0 did not preserve the DIS 5k validation dataset
149
150 ![Example 1](BEN2_demo_pictures/grid_example1.png)
151 ![Example 2](BEN2_demo_pictures/grid_example2.png)
152 ![Example 3](BEN2_demo_pictures/grid_example3.png)
153 ![Example 6](BEN2_demo_pictures/grid_example6.png)
154 ![Example 7](BEN2_demo_pictures/grid_example7.png)