Update app.py

app.py

@@ -549,6 +549,7 @@ def apply_mask(d, b):
     return masks[frame_selected], depths, frames
 
 def draw_mask(l, t, th, v, d, evt: gr.EventData):
+    global frames
     global depths
     global params
     global frame_selected
@@ -562,6 +563,8 @@ def draw_mask(l, t, th, v, d, evt: gr.EventData):
 
     if len(edge) == 0 or params["fnum"] != frame_selected or params["l"] != l:
       if len(edge) > 0:
+          clr = cv2.cvtColor(cv2.imread(frames[frame_selected]).astype(np.uint8), cv2.COLOR_BGR2HSV)
+          clr = cb2.medianBlur(clr, th)
           d["background"] = cv2.imread(depths[frame_selected]).astype(np.uint8)
           
           if d["background"].shape[0] == 2048: #height
@@ -571,19 +574,16 @@ def draw_mask(l, t, th, v, d, evt: gr.EventData):
           else:
             gradient = cv2.imread('./gradient_small.png').astype(np.uint8)
 
-      bg = cv2.cvtColor(d["background"], cv2.COLOR_RGBA2GRAY)
+      bg = cv2.cvtColor(d["background"], cv2.COLOR_RGBA2HSV)
+      bg[:,:,0] = clr[:,:,0]
 
-      diff = np.abs(bg.astype(np.int16)-cv2.cvtColor(gradient, cv2.COLOR_RGBA2GRAY).astype(np.int16)).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)
 
-      #bd = cv2.convertScaleAbs(cv2.Laplacian(bg, cv2.CV_64F, 5))
-        
-      bg = ((bg - 127) * 2).astype(np.uint8)
-
       #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
@@ -601,13 +601,13 @@ def draw_mask(l, t, th, v, d, evt: gr.EventData):
       res = center[label.flatten()]
       bg = res.reshape((bg.shape))
 
-      bg[mask>0] = 0
-      bg[bg==255] = 0
+      bg[mask>0] = (0,0,0)
+      bg[bg[:,:,2]==255] = (0,0,0)
 
       params["fnum"] = frame_selected
       params["l"] = l
 
-      d["layers"][0] = cv2.cvtColor(bg, cv2.COLOR_GRAY2RGBA)
+      d["layers"][0] = cv2.cvtColor(bg, cv2.COLOR_HSV2RGBA)
       edge = bg.copy()
     else:
       bg = edge.copy()