fkunn1326 commited on
Commit
6a97974
β€’
1 Parent(s): 0c43c79

Fix for cpu

Browse files
Files changed (1) hide show
  1. 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='cuda:0') # device can also be 'cpu'
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='cuda:0')
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