Update app.py

app.py

@@ -194,14 +194,6 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits', remove_bg
                     n = 0 #n = count-int(cframes/2)
                     depth_color_bg = cv2.imread(f"f{n}_dmap.png").astype(np.uint8)
                     raw_frame_bg = cv2.imread(f"f{n}.png").astype(np.uint8)
-
-                    dc = np.float32(depth_color.reshape((-1,3)))
-                    criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 4, 1.0)
-                    K = 16
-                    ret,label,center=cv2.kmeans(dc,K,None,criteria,4,cv2.KMEANS_PP_CENTERS)
-                    center = np.uint8(center)
-                    res = center[label.flatten()]
-                    depth_color = res.reshape((depth_color.shape))
                     
                     diff_d = np.abs(depth_color.astype(np.int16)-depth_color_bg.astype(np.int16))
                     diff_c = np.abs(raw_frame.astype(np.int16)-raw_frame_bg.astype(np.int16))
@@ -581,33 +573,36 @@ def draw_mask(l, t, v, d, evt: gr.EventData):
     pts = np.array(points, np.int32)
     pts = pts.reshape((-1,1,2))
 
-    scale = 1
-    delta = 0
-    ddepth = cv2.CV_16S
-
     if len(edge) == 0 or params["fnum"] != frame_selected or params["l"] != l:
       if len(edge) > 0:
           d["background"] = cv2.imread(depths[frame_selected]).astype(np.uint8)
 
       bg = cv2.cvtColor(d["background"], cv2.COLOR_RGBA2GRAY)
 
-      diff = (bg-cv2.cvtColor(gradient, cv2.COLOR_RGBA2GRAY)).astype(np.uint8)
-      mask = cv2.inRange(diff, -t, t)
+      diff = np.abs(bg.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)
-      grad = cv2.convertScaleAbs(cv2.Sobel(mask, ddepth, 1, 1, ksize=15-(t*2), scale=scale, delta=delta, borderType=cv2.BORDER_DEFAULT))
-
-      mask = mask + cv2.inRange(grad, 1, 255)
-
-      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]
-
+      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
+      #bg = cv2.convertScaleAbs(palette[bg]).astype(np.uint8) # Converting image back to uint
+
+      bg = np.float32(bg.reshape((-1,3)))
+      criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 4, 1.0)
+      ret,label,center=cv2.kmeans(bg,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
       bg[bg==255] = 0