Skip to content
Draft
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
51 changes: 43 additions & 8 deletions deeplabcut/generate_training_dataset/frame_extraction.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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):
Expand All @@ -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:
Expand Down
49 changes: 31 additions & 18 deletions deeplabcut/gui/tabs/extract_frames.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
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