Add Remove Background checkbox to Multi-Object tab with overlay option

app.py

@@ -386,7 +386,7 @@ def segment_video_simple(video_file, point_x, point_y, frame_idx, remove_bg):
 
 
 @spaces.GPU
-def segment_video_multi(video_file, objects_json):
+def segment_video_multi(video_file, objects_json, remove_bg=True):
     """
     Multi-object video segmentation.
     
@@ -397,6 +397,7 @@ def segment_video_multi(video_file, objects_json):
                 {"id": "ball", "init_frame": 0, "point_x": 360, "point_y": 640},
                 {"id": "player", "init_frame": 0, "point_x": 320, "point_y": 500}
             ]
+        remove_bg: If True, remove background. If False, overlay masks on original video.
     
     Returns:
         Tuple: (video1_path, video2_path, status_json)
@@ -537,11 +538,19 @@ def segment_video_multi(video_file, objects_json):
                             mask = cv2.resize(mask, (width, height), interpolation=cv2.INTER_NEAREST)
                         
                         mask_binary = (mask > 0.5).astype(np.uint8)
-                        
-                        # Remove background
-                        background = np.zeros_like(frame)
                         mask_3d = np.repeat(mask_binary[:, :, np.newaxis], 3, axis=2)
-                        frame = frame * mask_3d + background * (1 - mask_3d)
+                        
+                        if remove_bg:
+                            # Remove background - show only tracked object
+                            background = np.zeros_like(frame)
+                            frame = frame * mask_3d + background * (1 - mask_3d)
+                        else:
+                            # Overlay colored mask on original video
+                            # Create a colored overlay (e.g., semi-transparent green)
+                            overlay_color = np.array([0, 255, 0], dtype=np.uint8)  # Green
+                            overlay = np.ones_like(frame) * overlay_color
+                            alpha = 0.5  # Transparency
+                            frame = (frame * (1 - alpha * mask_3d) + overlay * alpha * mask_3d).astype(np.uint8)
                     
                     frame_bgr = cv2.cvtColor(frame.astype(np.uint8), cv2.COLOR_RGB2BGR)
                     out.write(frame_bgr)
@@ -740,6 +749,12 @@ def create_app():
                             placeholder='[{"id": "player", "init_frame": 165, "point_x": 180, "point_y": 320}]'
                         )
                         
+                        remove_bg_multi = gr.Checkbox(
+                            label="Remove Background", 
+                            value=True,
+                            info="If checked, shows only tracked objects. If unchecked, overlays masks on original video."
+                        )
+                        
                         process_btn_multi = gr.Button("🎬 Process Multi-Object", variant="primary")
                     
                     with gr.Column():
@@ -749,7 +764,7 @@ def create_app():
                 
                 process_btn_multi.click(
                     fn=segment_video_multi,
-                    inputs=[video_input_multi, objects_json],
+                    inputs=[video_input_multi, objects_json, remove_bg_multi],
                     outputs=[output_video_multi_1, output_video_multi_2, status_text_multi]
                 )