manual mask

app.py

@@ -117,7 +117,10 @@ with gr.Blocks(title="DiLightNet Demo") as demo:
                     seed=int(relighting_seed),
                     cfg=relighting_cfg
                 )
-                res = imageio.v3.imread(res_folder_path + '/relighting00.png')
+                mask_for_bg = imageio.v3.imread(res_folder_path + '/hint00_diffuse.png')[..., -1:] / 255.
+                res = imageio.v3.imread(res_folder_path + '/relighting00.png') / 255.
+                res = res * mask_for_bg  # + bg * (1. - mask_for_bg)
+                res = (res * 255).clip(0, 255).astype(np.uint8)
                 return res
 
 

demo/rm_bg.py

@@ -3,6 +3,7 @@ import rembg
 
 
 def rm_bg(img, use_sam=False):
+    print(type(img))
     output = rembg.remove(img)
     mask = np.array(output)[:, :, 3]