-
Notifications
You must be signed in to change notification settings - Fork 96
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add options to set background color when exporting video #1328
Changes from 5 commits
e5a8333
045ff9d
4762308
d8bdb43
2a82d3c
bd64c2a
330018c
b668551
5a71262
f1eef11
c7ec0a1
cd1663f
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change | ||||
---|---|---|---|---|---|---|
|
@@ -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,18 @@ 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 | ||||||
if background == "black": | ||||||
fill = 0 | ||||||
elif background == "grey": | ||||||
fill = 127 | ||||||
elif background == "white": | ||||||
# background == White | ||||||
fill = 255 | ||||||
else: | ||||||
raise ValueError(f"Invalid background color: {background}") | ||||||
video_frame_images = video_frame_images * 0 + fill | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This block of code is used to fill the frame with the specified color. However, it only supports three colors: black, grey, and white. If an unsupported color is passed, it raises a ValueError. Consider adding support for more colors or providing a better error message to guide users on what values are acceptable. - raise ValueError(f"Invalid background color: {background}")
+ raise ValueError(f"Invalid background color: {background}. Supported colors are 'black', 'grey', and 'white'.") |
||||||
|
||||||
if not loaded_chunk_idxs: | ||||||
print(f"No frames could be loaded from chunk {chunk_i}") | ||||||
|
@@ -497,6 +517,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 +536,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 | ||||||
shrivaths16 marked this conversation as resolved.
Show resolved
Hide resolved
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Match case sensitivity
Suggested change
|
||||||
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 +559,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 +717,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 +761,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}") | ||||||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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,44 @@ 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 --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() | ||
|
||
# Accumulate color 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. | ||
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 | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This test function checks if the video rendering with different background colors works as expected. It's a good practice to test with various inputs, and the use of pytest's parametrize decorator is appropriate here. However, the way you're checking the background color could be improved. Instead of splitting the frame into color channels and calculating the mean for each channel, you can calculate the mean color of the entire frame using - # Accumulate color channels
- b, g, r = cv2.split(frame)
- mean_b = np.mean(b)
- mean_g = np.mean(g)
- mean_r = np.mean(r)
+ # Calculate mean color of the frame
+ mean_color = np.mean(frame, axis=(0, 1)) |
||
|
||
|
||
def test_sleap_render(centered_pair_predictions): | ||
args = ( | ||
"-o testvis.avi -f 2 --scale 1.2 --frames 1,2 --video-index 0 " | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
These options should match case sensitivity, unless we handle this in visuals.py.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sure I'll make it lowercase in the docstrings and the option for GUI.
We could keep it in lowercase to be consistent.