Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 45 additions & 9 deletions deeplabcut/generate_training_dataset/frame_extraction.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,9 +9,28 @@
# Licensed under GNU Lesser General Public License v3.0
#

from typing import Collection, Iterable
from pathlib import Path


def _filter_config_videos(
configured_videos: Iterable[str | Path],
selected_videos: Collection[Path] | None,
) -> list[str | Path]:
"""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 = set(selected_videos)
return [video for video in configured_videos if 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
Expand Down Expand Up @@ -71,7 +90,7 @@ def extract_frames(
slider_width=25,
config3d=None,
extracted_cam=0,
videos_list=None,
videos_list: list[str | Path] | None = None,
):
"""Extracts frames from the project videos.

Expand Down Expand Up @@ -229,15 +248,23 @@ def extract_frames(
from skimage.util import img_as_ubyte

from deeplabcut.utils import auxiliaryfunctions, frameselectiontools

videos_list = (
None if videos_list is None else [Path(video) for video in videos_list]
)

config_file = Path(config)
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
Expand Down Expand Up @@ -407,7 +434,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):
Expand All @@ -427,9 +458,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:
Expand Down
48 changes: 30 additions & 18 deletions deeplabcut/gui/tabs/extract_frames.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,24 +51,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 Path(key) == 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
Expand Down Expand Up @@ -222,7 +230,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)
Expand Down
121 changes: 121 additions & 0 deletions tests/generate_training_dataset/test_frame_extraction.py
Original file line number Diff line number Diff line change
@@ -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
Loading