Update app.py

app.py

@@ -186,6 +186,23 @@ def make_video(video_path, outdir='./vis_video_depth', encoder='vits', remove_bg
 
             raw_frame = cv2.cvtColor(raw_frame, cv2.COLOR_BGR2BGRA)
             raw_frame[:, :, 3] = 255
+
+            if remove_bg == True:
+                if count >= int(len(cframes)/2):
+                    depth_gray_bg = cv2.imread(f"f{count-int(len(cframes)/2)}_mask.png").astype(np.uint8)
+                    raw_frame_bg = cv2.imread(f"f{count-int(len(cframes)/2)}.png").astype(np.uint8)
+                    
+                    mask_bg = cv2.inRange(np.abs(depth_gray-depth_gray_bg), 0, 16)
+                    mask_no_shadow = cv2.inRange(np.abs(raw_frame-raw_frame_bg), np.array([0,0,0,0]), np.array([16,16,16,0]))
+                    
+                    depth_gray[cv2.bitwise_and(mask_no_shadow, mask_bg)>0] = 0
+                    depth_color[cv2.bitwise_and(mask_no_shadow, mask_bg)>0] = (0,0,0)
+                    
+                    m = cv2.inRange(raw_frame, np.array([255,255,255,255]), np.array([255,255,255,255]))
+                    raw_frame[m>0] = (253,253,253,255)
+                    raw_frame[mask_bg>0] = (254,254,254,255)
+                    raw_frame[cv2.bitwise_and(mask_no_shadow, mask_bg)>0] = (255,255,255,255)
+            
             cv2.imwrite(f"f{count}.png", raw_frame)
             orig_frames.append(f"f{count}.png")
             
@@ -1281,28 +1298,22 @@ with gr.Blocks(css=css, js=js) as demo:
         locations = []
         avg = [0, 0]
         
-        if not coordinates:
-            locations = json.loads(example_coords)
-            for k, location in enumerate(locations):
+        locations = json.loads(coordinates)
+        for k, location in enumerate(locations):
+            if "tiles" in locations[k]:
+                locations[k]["heading"] = locations[k]["tiles"]["originHeading"]
+                locations[k]["pitch"] = locations[k]["tiles"]["originPitch"]
+            else:
+                locations[k]["heading"] = 0
+                locations[k]["pitch"] = 0
+
+            if "location" in locations[k]:
+                locations[k] = locations[k]["location"]["latLng"]
                 avg[0] = avg[0] + locations[k]["lat"]
                 avg[1] = avg[1] + locations[k]["lng"]
-        else:
-            locations = json.loads(coordinates)
-            for k, location in enumerate(locations):
-                if "tiles" in locations[k]:
-                    locations[k]["heading"] = locations[k]["tiles"]["originHeading"]
-                    locations[k]["pitch"] = locations[k]["tiles"]["originPitch"]
-                else:
-                    locations[k]["heading"] = 0
-                    locations[k]["pitch"] = 0
-
-                if "location" in locations[k]:
-                    locations[k] = locations[k]["location"]["latLng"]
-                    avg[0] = avg[0] + locations[k]["lat"]
-                    avg[1] = avg[1] + locations[k]["lng"]
-                else:
-                    locations[k]["lat"] = 0
-                    locations[k]["lng"] = 0
+            else:
+                locations[k]["lat"] = 0
+                locations[k]["lng"] = 0
                 
         if len(locations) > 0:
             avg[0] = avg[0] / len(locations)