diff --git a/deeplabcut/generate_training_dataset/frame_extraction.py b/deeplabcut/generate_training_dataset/frame_extraction.py index 7f7b22da6f..e4791d7c30 100755 --- a/deeplabcut/generate_training_dataset/frame_extraction.py +++ b/deeplabcut/generate_training_dataset/frame_extraction.py @@ -12,6 +12,28 @@ from pathlib import Path +def normalize_video_path(video: str | Path) -> Path: + return Path(video) + + +def _filter_config_videos( + configured_videos, + selected_videos, +) -> list: + """Return config video keys matching the selected video paths. + + The original config keys are returned so they remain valid for subsequent + config dictionary lookups. + """ + configured_videos = list(configured_videos) + + if selected_videos is None: + return configured_videos + + selected = {normalize_video_path(video) for video in selected_videos} + return [video for video in configured_videos if normalize_video_path(video) in selected] + + def select_cropping_area(config: str | Path, videos=None): """Interactively select the cropping area of all videos in the config. A user interface pops up with a frame to select the cropping parameters. Use the left click @@ -234,10 +256,14 @@ def extract_frames( cfg = auxiliaryfunctions.read_config(config_file) print("Config file read successfully.") - if videos_list is None: - videos = list(cfg.get("video_sets_original") or cfg["video_sets"]) - else: # filter video_list by the ones in the config file - videos = [v for v in cfg["video_sets"] if v in videos_list] + configured_videos = list(cfg.get("video_sets_original") or cfg["video_sets"]) + videos = _filter_config_videos(configured_videos, videos_list) + + if videos_list is not None and not videos: + raise ValueError( + "None of the selected videos matched the videos in the project " + "configuration. Selected videos may use a different path representation." + ) if mode == "manual": from deeplabcut.gui.widgets import launch_napari @@ -407,7 +433,11 @@ def extract_frames( else: # NO! has_failed.append(False) - if all(has_failed): + if not has_failed: + raise RuntimeError( + "No videos were processed. Check that the selected video paths match the entries in config.yaml" + ) + elif all(has_failed): print("Frame extraction failed. Video files must be corrupted.") return has_failed elif any(has_failed): @@ -427,9 +457,14 @@ def extract_frames( config_file = Path(config) cfg = auxiliaryfunctions.read_config(config_file) print("Config file read successfully.") - videos = sorted(cfg["video_sets"].keys()) - if videos_list is not None: # filter video_list by the ones in the config file - videos = [v for v in videos if v in videos_list] + + videos = _filter_config_videos(sorted(cfg["video_sets"]), videos_list) + if videos_list is not None and not videos: + raise ValueError( + "None of the selected videos matched the videos in the project " + "configuration. Selected videos may use a different path representation." + ) + project_path = Path(config).parents[0] labels_path = project_path / "labeled-data" try: diff --git a/deeplabcut/gui/tabs/extract_frames.py b/deeplabcut/gui/tabs/extract_frames.py index 6c04c7f43e..44f8d02327 100644 --- a/deeplabcut/gui/tabs/extract_frames.py +++ b/deeplabcut/gui/tabs/extract_frames.py @@ -15,6 +15,7 @@ from PySide6.QtCore import Qt from deeplabcut.generate_training_dataset import extract_frames +from deeplabcut.generate_training_dataset.frame_extraction import normalize_video_path from deeplabcut.gui.components import ( DefaultTab, VideoSelectionWidget, @@ -51,24 +52,32 @@ def select_cropping_area(config, videos=None): for video in videos: fc = FrameCropper(video) coords = fc.draw_bbox() - if coords: - temp = { - "crop": ", ".join( - map( - str, - [ - int(coords[0]), - int(coords[2]), - int(coords[1]), - int(coords[3]), - ], - ) + if not coords: + continue + + temp = { + "crop": ", ".join( + map( + str, + [ + int(coords[0]), + int(coords[2]), + int(coords[1]), + int(coords[3]), + ], ) - } - try: - cfg["video_sets"][video] = temp - except KeyError: - cfg["video_sets_original"][video] = temp + ) + } + + video_sets_name = "video_sets_original" if cfg.get("video_sets_original") else "video_sets" + video_sets = cfg[video_sets_name] + + matching_keys = [key for key in video_sets if normalize_video_path(key) == normalize_video_path(video)] + + if not matching_keys: + raise KeyError(f"Video is not present in the project configuration: {video}") + + video_sets[matching_keys[0]] = temp auxiliaryfunctions.write_config(config, cfg) return cfg @@ -222,7 +231,11 @@ def extract_frames(self): cluster_color=False, slider_width=slider_width, userfeedback=False, - videos_list=self.video_selection_widget.files or None, + videos_list=( + [str(video) for video in self.video_selection_widget.files] + if self.video_selection_widget.files + else None + ), ) self.worker, self.thread = move_to_separate_thread(func, capture_outputs=True) diff --git a/tests/generate_training_dataset/test_frame_extraction.py b/tests/generate_training_dataset/test_frame_extraction.py new file mode 100644 index 0000000000..5c906a32a9 --- /dev/null +++ b/tests/generate_training_dataset/test_frame_extraction.py @@ -0,0 +1,121 @@ +import os +from pathlib import Path + +import numpy as np +import pytest +from skimage import io + +from deeplabcut.generate_training_dataset.frame_extraction import _filter_config_videos, extract_frames +from deeplabcut.utils import auxfun_videos, auxiliaryfunctions, frameselectiontools + + +def test_extract_frames_accepts_path_videos_list(tmp_path, monkeypatch): + video = tmp_path / "videos" / "video.mp4" + video.parent.mkdir() + video.touch() + + cfg = { + "video_sets": { + str(video): {"crop": "0, 100, 0, 100"}, + }, + "numframes2pick": 1, + "start": 0.0, + "stop": 1.0, + } + + monkeypatch.setattr( + auxiliaryfunctions, + "read_config", + lambda _: cfg, + ) + + processed = [] + + class FakeVideoWriter: + def __init__(self, path): + processed.append(path) + + def __len__(self): + return 10 + + def set_to_frame(self, index): + pass + + def read_frame(self, crop=True): + return np.zeros((100, 100, 3), dtype=np.uint8) + + def close(self): + pass + + monkeypatch.setattr( + auxfun_videos, + "VideoWriter", + FakeVideoWriter, + ) + + # Isolate path filtering from frame-selection behavior. + monkeypatch.setattr( + frameselectiontools, + "UniformFramescv2", + lambda *args, **kwargs: [0], + ) + + # Avoid writing an actual PNG. + monkeypatch.setattr(io, "imsave", lambda *args, **kwargs: None) + + result = extract_frames( + tmp_path / "config.yaml", + mode="automatic", + algo="uniform", + videos_list=[video], + userfeedback=False, + ) + + assert processed == [str(video)] + assert result == [False] + + +class TestFilterConfigVideos: + def test_filter_config_videos_matches_path_to_string(self): + configured = [r"C:\project\videos\video.mp4"] + selected = [Path(r"C:\project\videos\video.mp4")] + + result = _filter_config_videos(configured, selected) + + assert result == configured + assert isinstance(result[0], str) + + def test_filter_config_videos_matches_string_to_path(self): + configured = [Path(r"C:\project\videos\video.mp4")] + selected = [r"C:\project\videos\video.mp4"] + + result = _filter_config_videos(configured, selected) + + assert result == configured + assert isinstance(result[0], Path) + + def test_filter_config_videos_preserves_original_config_key(self): + configured = [r"C:\project\videos\video.mp4"] + selected = [Path(r"C:\project\videos\video.mp4")] + + result = _filter_config_videos(configured, selected) + + assert result[0] is configured[0] + + def test_filter_config_videos_returns_all_when_selection_is_none(self): + configured = ["video-a.mp4", "video-b.mp4"] + + assert _filter_config_videos(configured, None) == configured + + def test_filter_config_videos_returns_empty_for_nonmatching_selection(self): + configured = ["video-a.mp4"] + selected = [Path("video-b.mp4")] + + assert _filter_config_videos(configured, selected) == [] + + @pytest.mark.skipif(os.name != "nt", reason="Windows path semantics") + def test_filter_config_videos_is_case_insensitive_on_windows(self): + configured = [r"C:\Project\Videos\VIDEO.MP4"] + selected = [Path(r"c:\project\videos\video.mp4")] + + assert _filter_config_videos(configured, selected) == configured