Spaces:
Running
on
Zero
Running
on
Zero
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)
|