ZhengPeng7
commited on
Commit
·
a0ef2a3
1
Parent(s):
0e5a7e4
Add General-lite as another option in the model selection.
Browse files
app.py
CHANGED
@@ -37,6 +37,7 @@ class ImagePreprocessor():
|
|
37 |
|
38 |
usage_to_weights_file = {
|
39 |
'General': 'BiRefNet',
|
|
|
40 |
'Portrait': 'BiRefNet-portrait',
|
41 |
'DIS': 'BiRefNet-DIS5K',
|
42 |
'HRSOD': 'BiRefNet-HRSOD',
|
@@ -55,8 +56,8 @@ def predict(image, resolution, weights_file):
|
|
55 |
global birefnet
|
56 |
# Load BiRefNet with chosen weights
|
57 |
_weights_file = '/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else usage_to_weights_file['General']))
|
58 |
-
print('
|
59 |
-
birefnet =
|
60 |
birefnet.to(device)
|
61 |
birefnet.eval()
|
62 |
|
|
|
37 |
|
38 |
usage_to_weights_file = {
|
39 |
'General': 'BiRefNet',
|
40 |
+
'General-lite': 'BiRefNet_T',
|
41 |
'Portrait': 'BiRefNet-portrait',
|
42 |
'DIS': 'BiRefNet-DIS5K',
|
43 |
'HRSOD': 'BiRefNet-HRSOD',
|
|
|
56 |
global birefnet
|
57 |
# Load BiRefNet with chosen weights
|
58 |
_weights_file = '/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else usage_to_weights_file['General']))
|
59 |
+
print('Using weights:', _weights_file)
|
60 |
+
birefnet = AutoModelForImageSegmentation.from_pretrained(_weights_file, trust_remote_code=True)
|
61 |
birefnet.to(device)
|
62 |
birefnet.eval()
|
63 |
|