ZhengPeng7 commited on
Commit
db1d921
·
1 Parent(s): 7403699

Update the usage names in the handler.

Browse files
Files changed (2) hide show
  1. handler.py +18 -7
  2. requirements.txt +1 -1
handler.py CHANGED
@@ -63,21 +63,24 @@ class ImagePreprocessor():
63
  usage_to_weights_file = {
64
  'General': 'BiRefNet',
65
  'General-HR': 'BiRefNet_HR',
66
- 'General-Lite': 'BiRefNet_lite',
67
- 'General-Lite-2K': 'BiRefNet_lite-2K',
68
- 'General-reso_512': 'BiRefNet-reso_512',
69
  'Matting': 'BiRefNet-matting',
70
- 'Matting-HR': 'BiRefNet_HR-Matting',
71
  'Portrait': 'BiRefNet-portrait',
 
 
 
 
72
  'DIS': 'BiRefNet-DIS5K',
73
  'HRSOD': 'BiRefNet-HRSOD',
74
  'COD': 'BiRefNet-COD',
75
  'DIS-TR_TEs': 'BiRefNet-DIS5K-TR_TEs',
76
- 'General-legacy': 'BiRefNet-legacy'
 
 
77
  }
78
 
79
  # Choose the version of BiRefNet here.
80
- usage = 'General'
81
 
82
  # Set resolution
83
  if usage in ['General-Lite-2K']:
@@ -87,7 +90,11 @@ elif usage in ['General-reso_512']:
87
  elif usage in ['General-HR', 'Matting-HR']:
88
  resolution = (2048, 2048)
89
  else:
90
- resolution = (1024, 1024)
 
 
 
 
91
 
92
  half_precision = True
93
 
@@ -123,6 +130,10 @@ class EndpointHandler():
123
 
124
  image = image_ori.convert('RGB')
125
  # Preprocess the image
 
 
 
 
126
  image_preprocessor = ImagePreprocessor(resolution=tuple(resolution))
127
  image_proc = image_preprocessor.proc(image)
128
  image_proc = image_proc.unsqueeze(0)
 
63
  usage_to_weights_file = {
64
  'General': 'BiRefNet',
65
  'General-HR': 'BiRefNet_HR',
66
+ 'Matting-HR': 'BiRefNet_HR-matting',
 
 
67
  'Matting': 'BiRefNet-matting',
 
68
  'Portrait': 'BiRefNet-portrait',
69
+ 'General-reso_512': 'BiRefNet_512x512',
70
+ 'General-Lite': 'BiRefNet_lite',
71
+ 'General-Lite-2K': 'BiRefNet_lite-2K',
72
+ # 'Anime-Lite': 'BiRefNet_lite-Anime',
73
  'DIS': 'BiRefNet-DIS5K',
74
  'HRSOD': 'BiRefNet-HRSOD',
75
  'COD': 'BiRefNet-COD',
76
  'DIS-TR_TEs': 'BiRefNet-DIS5K-TR_TEs',
77
+ 'General-legacy': 'BiRefNet-legacy',
78
+ 'General-dynamic': 'BiRefNet_dynamic',
79
+ 'Matting-dynamic': 'BiRefNet_dynamic-matting',
80
  }
81
 
82
  # Choose the version of BiRefNet here.
83
+ usage = 'General-dynamic'
84
 
85
  # Set resolution
86
  if usage in ['General-Lite-2K']:
 
90
  elif usage in ['General-HR', 'Matting-HR']:
91
  resolution = (2048, 2048)
92
  else:
93
+ if '-dynamic' in usage:
94
+ resolution = None
95
+ print('Using the original size (div by 32) for inference.')
96
+ else:
97
+ resolution = (1024, 1024)
98
 
99
  half_precision = True
100
 
 
130
 
131
  image = image_ori.convert('RGB')
132
  # Preprocess the image
133
+ if resolution is None:
134
+ resolution_div_by_32 = [int(int(reso)//32*32) for reso in image.size]
135
+ if resolution_div_by_32 != resolution:
136
+ resolution = resolution_div_by_32
137
  image_preprocessor = ImagePreprocessor(resolution=tuple(resolution))
138
  image_proc = image_preprocessor.proc(image)
139
  image_proc = image_proc.unsqueeze(0)
requirements.txt CHANGED
@@ -1,4 +1,4 @@
1
- torch==2.5.1
2
  torchvision
3
  numpy<2
4
  opencv-python
 
1
+ torch>=2.5.1
2
  torchvision
3
  numpy<2
4
  opencv-python