|
@@ -73,7 +73,7 @@ def save_result(alpha, path, im_path, trimap=None, fg_estimate=True, fg=None):
|
|
|
rgba = np.concatenate((fg, alpha), axis=-1)
|
|
|
cv2.imwrite(rgba_save_path, rgba)
|
|
|
|
|
|
- return fg
|
|
|
+ return fg, alpha_save_path, rgba_save_path
|
|
|
|
|
|
|
|
|
def reverse_transform(img, trans_info):
|
|
@@ -188,7 +188,7 @@ def predict(model,
|
|
|
|
|
|
save_path = os.path.join(save_dir, im_file)
|
|
|
mkdir(save_path)
|
|
|
- fg = save_result(
|
|
|
+ fg,alpha_save_path, rgba_save_path = save_result(
|
|
|
alpha,
|
|
|
save_path,
|
|
|
im_path=im_path,
|
|
@@ -214,4 +214,4 @@ def predict(model,
|
|
|
preprocess_cost_averager.reset()
|
|
|
infer_cost_averager.reset()
|
|
|
postprocess_cost_averager.reset()
|
|
|
- return alpha, fg
|
|
|
+ return alpha, fg, alpha_save_path, rgba_save_path
|