Visitar URL original
Niels/pcutoff arg (#2630) · DeepLabCut/DeepLabCut@11cedcc · GitHub
Skip to content

Commit 11cedcc

Browse files
authored
Niels/pcutoff arg (#2630)
1 parent 374025c commit 11cedcc

3 files changed

Lines changed: 51 additions & 16 deletions

File tree

‎deeplabcut/gui/tabs/create_videos.py‎

Lines changed: 20 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,24 @@ def _generate_layout_video_parameters(self, layout):
146146
)
147147
tmp_layout.addWidget(self.use_filtered_data_checkbox)
148148

149+
# Selector for p-cutoff
150+
pcutoff_widget = QtWidgets.QWidget()
151+
pcutoff_layout = _create_horizontal_layout(margins=(0, 0, 0, 0))
152+
pcutoff_label = QtWidgets.QLabel("Plotting confidence cutoff (pcutoff)")
153+
self.pcutoff_selector = QtWidgets.QDoubleSpinBox()
154+
self.pcutoff_selector.setMinimum(0.0)
155+
self.pcutoff_selector.setMaximum(1.0)
156+
self.pcutoff_selector.setValue(0.6)
157+
self.pcutoff_selector.setSingleStep(0.05)
158+
pcutoff_layout.addWidget(pcutoff_label)
159+
pcutoff_layout.addWidget(self.pcutoff_selector)
160+
pcutoff_widget.setLayout(pcutoff_layout)
161+
pcutoff_widget.setToolTip(
162+
"This value sets the confidence threshold, above which predictions are "
163+
"shown in the labeled videos."
164+
)
165+
tmp_layout.addWidget(pcutoff_widget)
166+
149167
# Plot trajectories
150168
self.plot_trajectories = QtWidgets.QCheckBox("Plot trajectories")
151169
self.plot_trajectories.setCheckState(Qt.Unchecked)
@@ -259,6 +277,7 @@ def create_videos(self):
259277
shuffle=shuffle,
260278
filtered=filtered,
261279
save_frames=self.create_high_quality_video.isChecked(),
280+
pcutoff=self.pcutoff_selector.value(),
262281
displayedbodyparts=bodyparts,
263282
draw_skeleton=self.draw_skeleton_checkbox.isChecked(),
264283
trailpoints=trailpoints,
@@ -273,22 +292,14 @@ def create_videos(self):
273292
failed_videos_str = ", ".join(failed_videos)
274293
self.root.writer.write(f"Failed to create videos from {failed_videos_str}.")
275294

276-
if all(videos_created):
277-
self.root.writer.write("Labeled videos created.")
278-
else:
279-
failed_videos = [
280-
video for success, video in zip(videos_created, videos) if not success
281-
]
282-
failed_videos_str = ", ".join(failed_videos)
283-
self.root.writer.write(f"Failed to create videos from {failed_videos_str}.")
284-
285295
if self.plot_trajectories.isChecked():
286296
deeplabcut.plot_trajectories(
287297
config=config,
288298
videos=videos,
289299
shuffle=shuffle,
290300
filtered=filtered,
291301
displayedbodyparts=bodyparts,
302+
pcutoff=self.pcutoff_selector.value(),
292303
)
293304

294305
def build_skeleton(self, *args):

‎deeplabcut/utils/make_labeled_video.py‎

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
Hao Wu, hwu01@g.harvard.edu contributed the original OpenCV class. Thanks!
2121
You can find the directory for your ffmpeg bindings by: "find / | grep ffmpeg" and then setting it.
2222
"""
23+
from __future__ import annotations
2324

2425
import argparse
2526
import os
@@ -382,7 +383,7 @@ def create_labeled_video(
382383
init_weights="",
383384
track_method="",
384385
superanimal_name="",
385-
pcutoff=0.6,
386+
pcutoff=None,
386387
skeleton=[],
387388
skeleton_color="white",
388389
dotsize=8,
@@ -501,6 +502,9 @@ def create_labeled_video(
501502
For multiple animals, must be either 'box', 'skeleton', or 'ellipse' and will
502503
be taken from the config.yaml file if none is given.
503504
505+
pcutoff: string, optional, default=None
506+
Overrides the pcutoff set in the project configuration to plot the trajectories.
507+
504508
overwrite: bool, optional, default=False
505509
If ``True`` overwrites existing labeled videos.
506510
@@ -560,13 +564,16 @@ def create_labeled_video(
560564
)
561565
"""
562566
if config == "":
563-
pass
567+
if pcutoff is None:
568+
pcutoff = 0.6
564569
else:
565570
cfg = auxiliaryfunctions.read_config(config)
566571
trainFraction = cfg["TrainingFraction"][trainingsetindex]
567572
track_method = auxfun_multianimal.get_track_method(
568573
cfg, track_method=track_method
569574
)
575+
if pcutoff is None:
576+
pcutoff = cfg["pcutoff"]
570577

571578
if init_weights == "":
572579
DLCscorer, DLCscorerlegacy = auxiliaryfunctions.get_scorer_name(
@@ -659,6 +666,7 @@ def create_labeled_video(
659666
keypoints_only,
660667
overwrite,
661668
init_weights=init_weights,
669+
pcutoff=pcutoff,
662670
confidence_to_alpha=confidence_to_alpha,
663671
)
664672

@@ -699,6 +707,7 @@ def proc_video(
699707
overwrite,
700708
video,
701709
init_weights="",
710+
pcutoff: float | None = None,
702711
confidence_to_alpha: Optional[Callable[[float], float]] = None,
703712
):
704713
"""Helper function for create_videos
@@ -716,6 +725,9 @@ def proc_video(
716725
if destfolder is None:
717726
destfolder = videofolder # where your folder with videos is.
718727

728+
if pcutoff is None:
729+
pcutoff = cfg["pcutoff"]
730+
719731
auxiliaryfunctions.attempt_to_make_folder(destfolder)
720732

721733
os.chdir(destfolder) # THE VIDEO IS STILL IN THE VIDEO FOLDER
@@ -751,7 +763,10 @@ def proc_video(
751763
s = "_id" if color_by == "individual" else "_bp"
752764
else:
753765
s = ""
754-
videooutname = filepath.replace(".h5", f"{s}_labeled.mp4")
766+
767+
videooutname = filepath.replace(
768+
".h5", f"{s}_p{int(100 * pcutoff)}_labeled.mp4"
769+
)
755770
if os.path.isfile(videooutname) and not overwrite:
756771
print("Labeled video already created. Skipping...")
757772
return
@@ -779,7 +794,7 @@ def proc_video(
779794
df,
780795
videooutname,
781796
inds,
782-
cfg["pcutoff"],
797+
pcutoff,
783798
cfg["dotsize"],
784799
cfg["alphavalue"],
785800
skeleton_color=skeleton_color,
@@ -801,7 +816,7 @@ def proc_video(
801816
cfg["dotsize"],
802817
cfg["colormap"],
803818
cfg["alphavalue"],
804-
cfg["pcutoff"],
819+
pcutoff,
805820
trailpoints,
806821
cropping,
807822
x1,
@@ -828,7 +843,7 @@ def proc_video(
828843
bbox=(x1, x2, y1, y2),
829844
codec=codec,
830845
output_path=videooutname,
831-
pcutoff=cfg["pcutoff"],
846+
pcutoff=pcutoff,
832847
dotsize=cfg["dotsize"],
833848
cmap=cfg["colormap"],
834849
color_by=color_by,

‎deeplabcut/utils/plotting.py‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
https://github.com/DeepLabCut/DeepLabCut/blob/master/AUTHORS
1818
Licensed under GNU Lesser General Public License v3.0
1919
"""
20+
from __future__ import annotations
2021

2122
import argparse
2223
import os
@@ -187,6 +188,7 @@ def plot_trajectories(
187188
resolution=100,
188189
linewidth=1.0,
189190
track_method="",
191+
pcutoff: float | None = None,
190192
):
191193
"""Plots the trajectories of various bodyparts across the video.
192194
@@ -251,6 +253,9 @@ def plot_trajectories(
251253
For multiple animals, must be either 'box', 'skeleton', or 'ellipse' and will
252254
be taken from the config.yaml file if none is given.
253255
256+
pcutoff: string, optional, default=None
257+
Overrides the pcutoff set in the project configuration to plot the trajectories.
258+
254259
Returns
255260
-------
256261
None
@@ -266,6 +271,10 @@ def plot_trajectories(
266271
)
267272
"""
268273
cfg = auxiliaryfunctions.read_config(config)
274+
275+
if pcutoff is None:
276+
pcutoff = cfg["pcutoff"]
277+
269278
track_method = auxfun_multianimal.get_track_method(cfg, track_method=track_method)
270279

271280
trainFraction = cfg["TrainingFraction"][trainingsetindex]
@@ -308,7 +317,7 @@ def plot_trajectories(
308317
linewidth,
309318
cfg["colormap"],
310319
cfg["alphavalue"],
311-
cfg["pcutoff"],
320+
pcutoff,
312321
suffix,
313322
imagetype,
314323
tmpfolder,

0 commit comments

Comments
 (0)