Commit
·
db1d921
1
Parent(s):
7403699
Update the usage names in the handler.
Browse files- handler.py +18 -7
- 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 |
-
'
|
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 |
-
|
|
|
|
|
|
|
|
|
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 |
torchvision
|
3 |
numpy<2
|
4 |
opencv-python
|
|
|
1 |
+
torch>=2.5.1
|
2 |
torchvision
|
3 |
numpy<2
|
4 |
opencv-python
|