Spaces:
Running
on
Zero
Running
on
Zero
ZhengPeng7
commited on
Commit
•
f70bf31
1
Parent(s):
7f03022
Add the option of the legacy general use weights without training on portrait seg data.
Browse files
app.py
CHANGED
@@ -45,7 +45,8 @@ usage_to_weights_file = {
|
|
45 |
'DIS': 'BiRefNet-DIS5K',
|
46 |
'HRSOD': 'BiRefNet-HRSOD',
|
47 |
'COD': 'BiRefNet-COD',
|
48 |
-
'DIS-TR_TEs': 'BiRefNet-DIS5K-TR_TEs'
|
|
|
49 |
}
|
50 |
|
51 |
birefnet = AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', usage_to_weights_file['General'])), trust_remote_code=True)
|
@@ -59,7 +60,7 @@ def predict(
|
|
59 |
resolution: str,
|
60 |
weights_file: Optional[str]
|
61 |
) -> Tuple[np.ndarray, np.ndarray]:
|
62 |
-
global birefnet
|
63 |
# Load BiRefNet with chosen weights
|
64 |
_weights_file = '/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else usage_to_weights_file['General']))
|
65 |
print('Using weights:', _weights_file)
|
|
|
45 |
'DIS': 'BiRefNet-DIS5K',
|
46 |
'HRSOD': 'BiRefNet-HRSOD',
|
47 |
'COD': 'BiRefNet-COD',
|
48 |
+
'DIS-TR_TEs': 'BiRefNet-DIS5K-TR_TEs',
|
49 |
+
'General-legacy': 'BiRefNet-legacy'
|
50 |
}
|
51 |
|
52 |
birefnet = AutoModelForImageSegmentation.from_pretrained('/'.join(('zhengpeng7', usage_to_weights_file['General'])), trust_remote_code=True)
|
|
|
60 |
resolution: str,
|
61 |
weights_file: Optional[str]
|
62 |
) -> Tuple[np.ndarray, np.ndarray]:
|
63 |
+
# global birefnet
|
64 |
# Load BiRefNet with chosen weights
|
65 |
_weights_file = '/'.join(('zhengpeng7', usage_to_weights_file[weights_file] if weights_file is not None else usage_to_weights_file['General']))
|
66 |
print('Using weights:', _weights_file)
|