Spaces:
Runtime error
Runtime error
Fix for cpu
Browse files- scripts/td_abg.py +2 -2
scripts/td_abg.py
CHANGED
@@ -58,7 +58,7 @@ def rmbg_fn(img):
|
|
58 |
|
59 |
def refinement(img, mask, fast, psp_L):
|
60 |
mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
|
61 |
-
refiner = refine.Refiner(device='
|
62 |
|
63 |
# Fast - Global step only.
|
64 |
# Smaller L -> Less memory usage; faster in fast mode.
|
@@ -106,7 +106,7 @@ def get_foreground(img, td_abg_enabled, h_split, v_split, n_cluster, alpha, th_r
|
|
106 |
if cascadePSP_enabled == True and td_abg_enabled == False:
|
107 |
mask = get_mask(img)
|
108 |
mask = (mask * 255).astype(np.uint8)
|
109 |
-
refiner = refine.Refiner(device='
|
110 |
mask = refiner.refine(img, mask, fast=fast, L=psp_L)
|
111 |
img = np.dstack((img, mask))
|
112 |
|
|
|
58 |
|
59 |
def refinement(img, mask, fast, psp_L):
|
60 |
mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
|
61 |
+
refiner = refine.Refiner(device='cpu') # device can also be 'cpu'
|
62 |
|
63 |
# Fast - Global step only.
|
64 |
# Smaller L -> Less memory usage; faster in fast mode.
|
|
|
106 |
if cascadePSP_enabled == True and td_abg_enabled == False:
|
107 |
mask = get_mask(img)
|
108 |
mask = (mask * 255).astype(np.uint8)
|
109 |
+
refiner = refine.Refiner(device='cpu')
|
110 |
mask = refiner.refine(img, mask, fast=fast, L=psp_L)
|
111 |
img = np.dstack((img, mask))
|
112 |
|