Update app.py

app.py

@@ -37,7 +37,6 @@ encoder2name = {
 
 blurin = "1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1"
 edge = []
-clrs = []
 gradient = None
 params = { "fnum":0, "l":16 }
 dcolor = []
@@ -226,6 +225,8 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits', blur_data
             cv2.imwrite(f"f{count}_dmap.png", depth_color)
             depth_frames.append(f"f{count}_dmap.png")
 
+            depth_gray = seg_frame(depth_gray)
+
             cv2.imwrite(f"f{count}_mask.png", depth_gray)
             masks.append(f"f{count}_mask.png")
             count += 1
@@ -414,16 +415,45 @@ def show_json(txt):
     return data[0]["video"]["path"], data[1]["path"], data[2], data[3]["background"]["path"], data[4], data[5]
 
 
+def seg_frame(newmask):
+
+    if depth.shape[0] == 2048: #height
+        gd = cv2.imread('./gradient_large.png', cv2.IMREAD_GRAYSCALE).astype(np.uint8)
+    elif depth.shape[0] == 1024:
+        gd = cv2.imread('./gradient.png', cv2.IMREAD_GRAYSCALE).astype(np.uint8)
+    else:
+        gd = cv2.imread('./gradient_small.png', cv2.IMREAD_GRAYSCALE).astype(np.uint8)
+   
+    newmask[np.absolute(newmask.astype(np.int16)-gd.astype(np.int16))<16] = 0
+    ret,newmask = cv2.threshold(newmask,0,255,cv2.THRESH_BINARY+cv2.THRESH_OTSU)
+
+    b = 1
+    d = 32
+    element = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * b + 1, 2 * b + 1), (b, b))
+    bd = cv2.erode(newmask, element)
+    element = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * d + 1, 2 * d + 1), (d, d))
+    bg = cv2.dilate(newmask, element)
+    bg[bg.shape[0]-64:bg.shape[0],0:bg.shape[1]] = 0
+
+    mask = np.zeros(depth.shape[:2],np.uint8)
+    # https://docs.opencv.org/4.x/d8/d83/tutorial_py_grabcut.html
+    # wherever it is marked white (sure foreground), change mask=1
+    # wherever it is marked black (sure background), change mask=0
+    mask[bg == 255] = 3
+    mask[bd == 255] = 1 #2: probable bg, 3: probable fg
+    
+    return mask
+
+
 def select_frame(d, evt: gr.SelectData):
     global dcolor
     global frame_selected
+    global depths
     global masks
     global edge
     
     if evt.index != frame_selected:
         edge = []
-        mask = cv2.imread(depths[frame_selected]).astype(np.uint8)
-        cv2.imwrite(masks[frame_selected], cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY))
         frame_selected = evt.index
 
     if len(dcolor) == 0:
@@ -431,7 +461,7 @@ def select_frame(d, evt: gr.SelectData):
     else:
         bg = "[" + str(dcolor[frame_selected])[1:-1] + ", 255]"
         
-    return masks[frame_selected], frame_selected, bg
+    return depths[frame_selected], frame_selected, bg
 
 def switch_rows(v):
     global frames
@@ -516,7 +546,7 @@ def reset_mask():
     cv2.imwrite(masks[frame_selected], cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY))
     return masks[frame_selected], depths
 
-def apply_mask(d, b):
+def apply_mask(d,b):
     global frames
     global frame_selected
     global masks
@@ -525,29 +555,20 @@ def apply_mask(d, b):
 
     edge = []
     mask = cv2.cvtColor(d["layers"][0], cv2.COLOR_RGBA2GRAY)
-    mask[mask<255] = 0
-    b = b*2+1
-    dilation = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * b + 1, 2 * b + 1), (b, b))
-    mask = cv2.dilate(mask, dilation)
-    mask_b = cv2.GaussianBlur(mask, (b,b), 0)
-    b = b*2+1
-    dilation = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * b + 1, 2 * b + 1), (b, b))
-    dmask = cv2.dilate(mask, dilation)
-    dmask_b = cv2.GaussianBlur(dmask, (b,b), 0)
-
-    for k, mk in enumerate(masks):
-        if k != frame_selected and k < len(depths):
-            cv2.imwrite(masks[k], dmask)
-            frame = cv2.imread(frames[k], cv2.IMREAD_UNCHANGED).astype(np.uint8)
-            frame[:, :, 3] = dmask_b
-            cv2.imwrite(frames[k], frame)
-        
+    mask = (mask - 128).astype(np.uint8)
+
+    mask, bgdModel, fgdModel = cv2.grabCut(d["background"], 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] = 255 - mask_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[mask==128] = 0
+    d["layers"][0] = cv2.cvtColor(mask, cv2.COLOR_GRAY2RGBA)
     
-    cv2.imwrite(masks[frame_selected], mask) #d["background"]
-    return masks[frame_selected], depths, frames
+    return gr.ImageEditor(value=d), depths, frames
 
 def draw_mask(l, t, b, v, d, evt: gr.EventData):
     global frames
@@ -557,63 +578,24 @@ def draw_mask(l, t, b, v, d, evt: gr.EventData):
     global masks
     global gradient
     global edge
-    global clrs
     
     points = json.loads(v)
     pts = np.array(points, np.int32)
     pts = pts.reshape((-1,1,2))
 
     if len(edge) == 0 or params["fnum"] != frame_selected or params["l"] != l:
-      clrs = cv2.cvtColor(cv2.imread(frames[frame_selected]).astype(np.uint8), cv2.COLOR_BGR2HSV)
-      clrs = cv2.medianBlur(clrs, b)
       if len(edge) > 0:
           d["background"] = cv2.imread(depths[frame_selected]).astype(np.uint8)
-          
-          if d["background"].shape[0] == 2048: #height
-            gradient = cv2.imread('./gradient_large.png').astype(np.uint8)
-          elif d["background"].shape[0] == 1024:
-            gradient = cv2.imread('./gradient.png').astype(np.uint8)
-          else:
-            gradient = cv2.imread('./gradient_small.png').astype(np.uint8)
+          mask = cv2.imread(masks[frame_selected], cv2.IMREAD_GRAYSCALE).astype(np.uint8) + 128
+          mask[mask==128] = 0
 
       bg = cv2.cvtColor(d["background"], cv2.COLOR_RGBA2GRAY)
-      #bg = cv2.cvtColor(bg, cv2.COLOR_RGB2HSV)
-      #bg[:,:,0] = clrs[:,:,0]
-      #bg = cv2.cvtColor(bg, cv2.COLOR_HSV2RGB)
-      bd = cv2.convertScaleAbs(cv2.Laplacian(bg, cv2.CV_64F, 7)).astype(np.uint8)
-
-      diff = np.abs(cv2.cvtColor(d["background"], cv2.COLOR_RGBA2GRAY).astype(np.int16)-cv2.cvtColor(gradient, cv2.COLOR_RGBA2GRAY).astype(np.int16)).astype(np.uint8)
-      mask = cv2.inRange(diff, 0, t)
-      #kernel = np.ones((c,c),np.float32)/(c*c)
-      #mask = cv2.filter2D(mask,-1,kernel)
-      dilation = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15-(t*2+1), 15-(t*2+1)), (t, t))
-      mask = cv2.dilate(mask, dilation)
-
-      #indices = np.arange(0,256)   # List of all colors 
-      #divider = np.linspace(0,255,l+1)[1] # we get a divider
-      #quantiz = np.intp(np.linspace(0,255,l)) # we get quantization colors
-      #color_levels = np.clip(np.intp(indices/divider),0,l-1) # color levels 0,1,2..
-      #palette = quantiz[color_levels]
-      
-      #for i in range(l):
-      #    bg[(bg >= i*255/l) & (bg < (i+1)*255/l)] = i*255/(l-1)
-      #bg = cv2.convertScaleAbs(palette[bg]).astype(np.uint8) # Converting image back to uint
-
-      res = np.float32(bg.reshape((-1,1)))
-      criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 4, 1.0)
-      ret,label,center=cv2.kmeans(res,l,None,criteria,4,cv2.KMEANS_PP_CENTERS)
-      center = np.uint8(center)
-      res = center[label.flatten()]
-      bg = res.reshape((bg.shape))
-
-      bg[bd>4] = 0
-      bg[mask>0] = 0
       bg[bg==255] = 0
 
       params["fnum"] = frame_selected
       params["l"] = l
 
-      d["layers"][0] = cv2.cvtColor(bg, cv2.COLOR_GRAY2RGBA)
+      d["layers"][0] = cv2.cvtColor(mask, cv2.COLOR_GRAY2RGBA)
       edge = bg.copy()
     else:
       bg = edge.copy()
@@ -621,12 +603,11 @@ def draw_mask(l, t, b, v, d, evt: gr.EventData):
     x = points[len(points)-1][0]
     y = points[len(points)-1][1]
 
-                                                #int(t*256/l)
-    mask = cv2.floodFill(bg, None, (x, y), 1, 0, 256, (4 | cv2.FLOODFILL_FIXED_RANGE))[2] #(4 | cv2.FLOODFILL_FIXED_RANGE | cv2.FLOODFILL_MASK_ONLY | 255 << 8)
+    sel = cv2.floodFill(d["layers"][0], None, (x, y), 1, 2, 2, (4 | cv2.FLOODFILL_FIXED_RANGE))[2] #(4 | cv2.FLOODFILL_FIXED_RANGE | cv2.FLOODFILL_MASK_ONLY | 255 << 8)
     # 255 << 8 tells to fill with the value 255)
-    mask = mask[1:mask.shape[0]-1, 1:mask.shape[1]-1]
-    
-    d["layers"][0][mask>0] = (255,255,255,255)
+    sel = sel[1:sel.shape[0]-1, 1:sel.shape[1]-1]
+
+    d["layers"][0][sel==0] = (0,0,0,0)
     
     return gr.ImageEditor(value=d)