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
Files changed (1) hide show
  1. app.py +3 -2
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)