transparent background

main.py

@@ -129,7 +129,7 @@ def save_temp_imgs(imgs):
 
 def getThumbnails(imgs):
     thumbs = []
-    thumb_size = (256, 256)
+    thumb_size = (512, 512)
     for img in imgs:
         th = img.copy()
         th.thumbnail(thumb_size)
@@ -248,14 +248,19 @@ def extract_frames(video_path, fps):
     return frames
 
 
-def removebg(img, rembg_session):
-    result = Image.new("RGB", img.size, "#ffffff")
+def removebg(img, rembg_session, transparent=False):
+    
+    if transparent:
+        result = Image.new('RGBA', img.size, (0, 0, 0, 0))
+    else:
+        result = Image.new("RGB", img.size, "#ffffff")
     out = rembg.remove(img, session=rembg_session)
     result.paste(out, mask=out)
     return result
 
 
 def prepare_inputs_train(images, bg_remove, dwpose, rembg_session):
+    print("remove background", bg_remove)
     if bg_remove:
         images = [removebg(img, rembg_session) for img in images]
 
@@ -270,7 +275,7 @@ def prepare_inputs_train(images, bg_remove, dwpose, rembg_session):
     return in_img, in_pose, train_imgs, train_poses
         
 
-def prepare_inputs_inference(in_img, in_vid, fps, dwpose, resize='target', is_app=False):
+def prepare_inputs_inference(in_img, in_vid, fps, dwpose, rembg_session, bg_remove, resize='target', is_app=False):
     progress=gr.Progress(track_tqdm=True)
 
     print("prepare_inputs_inference")
@@ -278,7 +283,10 @@ def prepare_inputs_inference(in_img, in_vid, fps, dwpose, resize='target', is_ap
     in_pose = get_pose(in_img, dwpose, "in_pose.png")
     
     frames = extract_frames(in_vid, fps)
-    #frames = [removebg(img, rembg_session) for img in frames]
+    print("remove background", bg_remove)
+    if bg_remove:
+        in_img = removebg(in_img, rembg_session)
+        #frames = [removebg(img, rembg_session) for img in frames]
     if debug:
         for i, frame in enumerate(frames):
             frame.save("out/frame_"+str(i)+".png")
@@ -317,14 +325,14 @@ def prepare_inputs_inference(in_img, in_vid, fps, dwpose, resize='target', is_ap
             tpose.save("out/"+"tar_pose"+str(len(target_poses_cropped))+".png")
         target_poses_cropped.append(tpose)
         
-    return target_poses_cropped, in_pose
+    return in_img, target_poses_cropped, in_pose
 
 
 def prepare_inputs(images, in_vid, fps, bg_remove, dwpose, rembg_session, resize='target', is_app=False):
     
     in_img, in_pose, train_imgs, train_poses = prepare_inputs_train(images, bg_remove, dwpose, rembg_session)
     
-    target_poses_cropped, _ = prepare_inputs_inference(in_img, in_vid, fps, dwpose, resize, is_app)
+    in_img, target_poses_cropped, _ = prepare_inputs_inference(in_img, in_vid, fps, dwpose, rembg_session, bg_remove, resize, is_app)
     
     
     return in_img, in_pose, train_imgs, train_poses, target_poses_cropped
@@ -1110,7 +1118,7 @@ def run_inference(images, video_path, train_steps=100, inference_steps=10, fps=1
     images = [img[0] for img in images]
     in_img = images[0]
 
-    target_poses, in_pose = prepare_inputs_inference(in_img, video_path, fps, dwpose, 'target', is_app)
+    in_img, target_poses, in_pose = prepare_inputs_inference(in_img, video_path, fps, dwpose, rembg_session, bg_remove, 'target', is_app)
     
     results = inference(modelId, in_img, in_pose, target_poses, inference_steps, None, vae, unet, image_encoder_p, is_app)
     #urls = save_temp_imgs(results)
@@ -1121,6 +1129,9 @@ def run_inference(images, video_path, train_steps=100, inference_steps=10, fps=1
         else:
             gen_vid(results, out_vid+'.webm', fps, 'webm')
     
+    
+    results = [removebg(img, rembg_session, True) for img in results]
+    
     print("Done!")
     
     return out_vid+'.webm', results, getThumbnails(results)