Update app.py

app.py

@@ -134,7 +134,7 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits'):
             if not ret:
                 break
 
-            frame = cv2.cvtColor(raw_frame, cv2.COLOR_BGR2RGB) / 255.0
+            frame = cv2.cvtColor(raw_frame, cv2.COLOR_BGR2RGBA) / 255.0
             frame_pil = Image.fromarray((frame * 255).astype(np.uint8))
             frame = transform({'image': frame})['image']
             
@@ -464,6 +464,7 @@ def reset_mask():
     return masks[frame_selected], f"f{frame_selected}_dmap.png"
 
 def apply_mask(d, b):
+    global frames
     global frame_selected
     global masks
     global depths
@@ -477,17 +478,23 @@ def apply_mask(d, b):
     dilation = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * b + 1, 2 * b + 1), (b, b))
     dmask = cv2.dilate(mask, dilation)
 
-    d["background"][mask==255] = (255,255,255,255)
-    mask = cv2.inRange(d["background"], np.array([255,255,255,255]), np.array([255,255,255,255]))
+    d["background"][mask==255] = (0,0,0,0)
+    mask = cv2.inRange(d["background"], np.array([0,0,0,0]), np.array([0,0,0,0]))
 
     for k, mk in enumerate(masks):
         if k != frame_selected and k < len(depths):
             masks[k] = cv2.imread(depths[k]).astype(np.uint8)
-            masks[k][dmask<255] = (255,255,255)
+            masks[k][dmask<255] = (0,0,0)
 
+    ksize = np.sqrt(np.sum(mask) / 255)
+    frame = cv2.imread(frames[frame_selected]).astype(np.uint8)
+    dmask = cv2.GaussianBlur(dmask, (ksize,ksize), 0)
+    frame[3] = 255 - dmask
+    cv2.imwrite(frames[frame_selected], frame)
+    
     masks[frame_selected] = d["background"]
     cv2.imwrite(f"f{frame_selected}_dmap.png", masks[frame_selected])
-    return masks[frame_selected], f"f{frame_selected}_dmap.png"
+    return masks[frame_selected], f"f{frame_selected}_dmap.png", frames
 
 def draw_mask(l, t, v, d, evt: gr.EventData):
     global params
@@ -786,7 +793,7 @@ with gr.Blocks(css=css, js=js) as demo:
               bsize = gr.Slider(label="Border size", value=15, maximum=256, minimum=1, step=2)
               mouse = gr.Textbox(elem_id="mouse", value="""[]""", interactive=False)
               mouse.input(fn=draw_mask, show_progress="minimal", inputs=[levels, tolerance, mouse, output_mask], outputs=[output_mask])
-              apply.click(fn=apply_mask, inputs=[output_mask, bsize], outputs=[output_mask, depth_file])
+              apply.click(fn=apply_mask, inputs=[output_mask, bsize], outputs=[output_mask, depth_file, output_frame])
               reset.click(fn=reset_mask, inputs=None, outputs=[output_mask, depth_file])
         with gr.Column():
             model_type = gr.Dropdown([("small", "vits"), ("base", "vitb"), ("large", "vitl")], type="value", value="vits", label='Model Type')