Add options to set background color when exporting video

docs/guides/cli.md

@@ -389,6 +389,9 @@ optional arguments:
   --distinctly_color DISTINCTLY_COLOR
                         Specify how to color instances. Options include: "instances",
                         "edges", and "nodes" (default: "instances")
+  --background BACKGROUND
+                        Specify the type of background to be used to save the videos.
+                        Options: original, black, white and grey. (default: "original")
 ```
 
 ## Debugging

sleap/config/labeled_clip_form.yaml

@@ -18,6 +18,10 @@ main:
     label: Use GUI Visual Settings (colors, line widths)
     type: bool
     default: true
+  - name: background
+    label: Video Background
+    type: list
+    options: original,black,white,grey
   - name: open_when_done
     label: Open When Done Saving
     type: bool

sleap/gui/commands.py

@@ -1295,6 +1295,7 @@ def do_action(context: CommandContext, params: dict):
             frames=list(params["frames"]),
             fps=params["fps"],
             color_manager=params["color_manager"],
+            background=params["background"],
             show_edges=params["show edges"],
             edge_is_wedge=params["edge_is_wedge"],
             marker_size=params["marker size"],
@@ -1354,6 +1355,7 @@ def ask(context: CommandContext, params: dict) -> bool:
         params["fps"] = export_options["fps"]
         params["scale"] = export_options["scale"]
         params["open_when_done"] = export_options["open_when_done"]
+        params["background"] = export_options["background"]
 
         params["crop"] = None
 

sleap/io/video.py

@@ -1118,8 +1118,9 @@ def get_frames(self, idxs: Union[int, Iterable[int]]) -> np.ndarray:
 
     def get_frames_safely(self, idxs: Iterable[int]) -> Tuple[List[int], np.ndarray]:
         """Return list of frame indices and frames which were successfully loaded.
+        Args:
+            idxs: An iterable object that contains the indices of frames.
 
-        idxs: An iterable object that contains the indices of frames.
 
         Returns: A tuple of (frame indices, frames), where
             * frame indices is a subset of the specified idxs, and

sleap/io/visuals.py

@@ -27,7 +27,13 @@
 _sentinel = object()
 
 
-def reader(out_q: Queue, video: Video, frames: List[int], scale: float = 1.0):
+def reader(
+    out_q: Queue,
+    video: Video,
+    frames: List[int],
+    scale: float = 1.0,
+    background: str = "original",
+):
     """Read frame images from video and send them into queue.
 
     Args:
@@ -36,11 +42,13 @@ def reader(out_q: Queue, video: Video, frames: List[int], scale: float = 1.0):
         video: The `Video` object to read.
         frames: Full list frame indexes we want to read.
         scale: Output scale for frame images.
+        background: output video background. Either original, black, white, grey
 
     Returns:
         None.
     """
 
+    background = background.lower()
     cv2.setNumThreads(usable_cpu_count())
 
     total_count = len(frames)
@@ -64,6 +72,16 @@ def reader(out_q: Queue, video: Video, frames: List[int], scale: float = 1.0):
             loaded_chunk_idxs, video_frame_images = video.get_frames_safely(
                 frames_idx_chunk
             )
+            if background != "original":
+                # fill the frame with the color
+                fill_values = {"black": 0, "grey": 127, "white": 255}
+                try:
+                    fill = fill_values[background]
+                except KeyError:
+                    raise ValueError(
+                        f"Invalid background color: {background}. Options include: {', '.join(fill_values.keys())}"
+                    )
+                video_frame_images = video_frame_images * 0 + fill
 
             if not loaded_chunk_idxs:
                 print(f"No frames could be loaded from chunk {chunk_i}")
@@ -497,6 +515,7 @@ def save_labeled_video(
     fps: int = 15,
     scale: float = 1.0,
     crop_size_xy: Optional[Tuple[int, int]] = None,
+    background: str = "original",
     show_edges: bool = True,
     edge_is_wedge: bool = False,
     marker_size: int = 4,
@@ -515,6 +534,7 @@ def save_labeled_video(
         fps: Frames per second for output video.
         scale: scale of image (so we can scale point locations to match)
         crop_size_xy: size of crop around instances, or None for full images
+        background: output video background. Either original, black, white, grey
         show_edges: whether to draw lines between nodes
         edge_is_wedge: whether to draw edges as wedges (draw as line if False)
         marker_size: Size of marker in pixels before scaling by `scale`
@@ -537,7 +557,7 @@ def save_labeled_video(
     q2 = Queue(maxsize=10)
     progress_queue = Queue()
 
-    thread_read = Thread(target=reader, args=(q1, video, frames, scale))
+    thread_read = Thread(target=reader, args=(q1, video, frames, scale, background))
     thread_mark = VideoMarkerThread(
         in_q=q1,
         out_q=q2,
@@ -695,6 +715,15 @@ def main(args: list = None):
             "and 'nodes' (default: 'nodes')"
         ),
     )
+    parser.add_argument(
+        "--background",
+        type=str,
+        default="original",
+        help=(
+            "Specify the type of background to be used to save the videos."
+            "Options for background: original, black, white and grey"
+        ),
+    )
     args = parser.parse_args(args=args)
     labels = Labels.load_file(
         args.data_path, video_search=[os.path.dirname(args.data_path)]
@@ -730,6 +759,7 @@ def main(args: list = None):
         marker_size=args.marker_size,
         palette=args.palette,
         distinctly_color=args.distinctly_color,
+        background=args.background,
     )
 
     print(f"Video saved as: {filename}")

tests/io/test_visuals.py

@@ -1,6 +1,7 @@
 import numpy as np
 import os
 import pytest
+import cv2
 from sleap.io.dataset import Labels
 from sleap.io.visuals import (
     save_labeled_video,
@@ -63,6 +64,46 @@ def test_serial_pipeline(centered_pair_predictions, tmpdir):
     )
 
 
+@pytest.mark.parametrize("background", ["original", "black", "white", "grey"])
+def test_sleap_render_with_different_backgrounds(background):
+    args = (
+        f"-o test_{background}.avi -f 2 --scale 1.2 --frames 1,2 --video-index 0 "
+        f"--background {background} "
+        "tests/data/json_format_v2/centered_pair_predictions.json".split()
+    )
+    sleap_render(args)
+    assert (
+        os.path.exists(f"test_{background}.avi")
+        and os.path.getsize(f"test_{background}.avi") > 0
+    )
+
+    # Check if the background is set correctly if not original background
+    if background != "original":
+        saved_video_path = f"test_{background}.avi"
+        cap = cv2.VideoCapture(saved_video_path)
+        ret, frame = cap.read()
+
+        # Calculate mean color of the channels
+        b, g, r = cv2.split(frame)
+        mean_b = np.mean(b)
+        mean_g = np.mean(g)
+        mean_r = np.mean(r)
+
+        # Set threshold values. Color is white if greater than white threshold, black
+        # if less than grey threshold and grey if in between both threshold values.
+        white_threshold = 240
+        grey_threshold = 40
+
+        # Check if the average color is white, grey, or black
+        if all(val > white_threshold for val in [mean_b, mean_g, mean_r]):
+            background_color = "white"
+        elif all(val < grey_threshold for val in [mean_b, mean_g, mean_r]):
+            background_color = "black"
+        else:
+            background_color = "grey"
+        assert background_color == background
+
+
 def test_sleap_render(centered_pair_predictions):
     args = (
         "-o testvis.avi -f 2 --scale 1.2 --frames 1,2 --video-index 0 "