Update app.py

app.py

@@ -576,10 +576,11 @@ def draw_mask(l, t, b, v, d, evt: gr.EventData):
           else:
             gradient = cv2.imread('./gradient_small.png').astype(np.uint8)
 
-      bg = cv2.cvtColor(d["background"], cv2.COLOR_RGBA2RGB)
-      bg = cv2.cvtColor(bg, cv2.COLOR_RGB2HSV)
-      bg[:,:,0] = clrs[:,:,0]
-      bg = cv2.cvtColor(bg, cv2.COLOR_HSV2RGB)
+      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, 5)).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)
@@ -598,20 +599,21 @@ def draw_mask(l, t, b, v, d, evt: gr.EventData):
       #    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,3)))
+      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[mask>0] = (0,0,0)
-      bg[bg[:,:,2]==255] = (0,0,0)
+      bg[bd>b] = 0
+      bg[mask>0] = 0
+      bg[bg==255] = 0
 
       params["fnum"] = frame_selected
       params["l"] = l
 
-      d["layers"][0] = cv2.cvtColor(bg, cv2.COLOR_RGB2RGBA)
+      d["layers"][0] = cv2.cvtColor(bg, cv2.COLOR_GRAY2RGBA)
       edge = bg.copy()
     else:
       bg = edge.copy()