Update app.py

app.py

@@ -567,7 +567,7 @@ def apply_mask(d,b):
     frame[:, :, 3] = mask * 255
     cv2.imwrite(frames[frame_selected], frame)
 
-    mask = cv2.imread(masks[frame_selected], cv2.IMREAD_GRAYSCALE).astype(np.uint8) + 128
+    mask = cv2.imread(masks[frame_selected], cv2.IMREAD_GRAYSCALE).astype(np.uint8)
     mask[mask==128] = 0
     d["layers"][0] = cv2.cvtColor(mask, cv2.COLOR_GRAY2RGBA)
     
@@ -613,6 +613,17 @@ def draw_mask(l, t, b, v, d, evt: gr.EventData):
     sel = sel[1:sel.shape[0]-1, 1:sel.shape[1]-1]
 
     d["layers"][0][sel==0] = (0,0,0,0)
+
+
+    mask = cv2.cvtColor(d["layers"][0], cv2.COLOR_RGBA2GRAY) - 128
+    mask[mask<0] = 0
+    print(mask[mask>0])
+    mask, bgdModel, fgdModel = cv2.grabCut(cv2.cvtColor(d["background"], cv2.COLOR_RGBA2RGB), mask, None,None,None,15, cv2.GC_INIT_WITH_MASK)
+    mask = np.where((mask==2)|(mask==0),1,0).astype('uint8')
+
+    frame = cv2.imread(frames[frame_selected], cv2.IMREAD_UNCHANGED).astype(np.uint8)
+    frame[:, :, 3] = mask * 255
+    cv2.imwrite(frames[frame_selected], frame)
     
     return gr.ImageEditor(value=d)