Visitar URL original
Improve GUI Modelzoo Tab by maximpavliv · Pull Request #3038 · DeepLabCut/DeepLabCut · GitHub
Skip to content

Improve GUI Modelzoo Tab #3038

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.

Sign up for GitHub

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

Closed
maximpavliv wants to merge 8 commits into from
Closed

Improve GUI Modelzoo Tab #3038

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
ab69e7a
video_inference_superanimal: add create_labeled_video argument
maximpavliv Jun 30, 2025
ecf9a13
Rename method
maximpavliv Jul 1, 2025
51cec0b
Add create labeled video checkbox
maximpavliv Jul 1, 2025
ae6b39a
Show/hide video adaptation parameter widgets
maximpavliv Jul 3, 2025
88a5793
Use set_combo_items() to update combo boxes
maximpavliv Jul 3, 2025
6ad0133
Add batch size combo boxes for pose model and detector
maximpavliv Jul 3, 2025
02f6439
Pass batch sizes to video_inference_superanimal()
maximpavliv Jul 3, 2025
7e22bac
Black
maximpavliv Jul 3, 2025
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
162 changes: 119 additions & 43 deletions deeplabcut/gui/tabs/modelzoo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
_create_label_widget,
DefaultTab,
VideoSelectionWidget,
set_combo_items,
)
from deeplabcut.gui.utils import move_to_separate_thread
from deeplabcut.gui.widgets import ClickableLabel
Expand Down Expand Up @@ -64,7 +65,7 @@ def _set_page(self):
self._build_torch_attributes()

self.run_button = QtWidgets.QPushButton("Run")
self.run_button.clicked.connect(self.run_video_adaptation)
self.run_button.clicked.connect(self.run_video_inference_superanimal)
self.main_layout.addWidget(self.run_button, alignment=Qt.AlignRight)

self.home_button = QtWidgets.QPushButton("Return to Welcome page")
Expand Down Expand Up @@ -116,16 +117,27 @@ def _build_common_attributes(self) -> None:
)
action.triggered.connect(self.select_folder)

self.create_labeled_video_checkbox = QtWidgets.QCheckBox("Create labeled video")
self.create_labeled_video_checkbox.setChecked(True)

batch_size_combo_label = QtWidgets.QLabel("Pose model batch size")
self.batch_size_combo = QtWidgets.QComboBox()
self.batch_size_combo.setMaximumWidth(100)
self.batch_size_combo.addItems([str(2**i) for i in range(6)])
self.batch_size_combo.setCurrentIndex(0)

settings_layout.addWidget(section_title, 0, 0)
settings_layout.addWidget(model_combo_text, 1, 0)
settings_layout.addWidget(self.model_combo, 1, 1)
settings_layout.addWidget(net_type_text, 2, 0)
settings_layout.addWidget(self.net_type_selector, 2, 1)
settings_layout.addWidget(self.detector_type_text, 3, 0)
settings_layout.addWidget(self.detector_type_selector, 3, 1)

settings_layout.addWidget(loc_label, 4, 0)
settings_layout.addWidget(self.loc_line, 4, 1)
settings_layout.addWidget(self.create_labeled_video_checkbox, 5, 0)
settings_layout.addWidget(batch_size_combo_label, 6, 0)
settings_layout.addWidget(self.batch_size_combo, 6, 1)

self.settings_widget = QtWidgets.QWidget()
self.settings_widget.setLayout(settings_layout)
Expand Down Expand Up @@ -160,7 +172,9 @@ def _build_tf_attributes(self) -> None:
self.adapt_checkbox = QtWidgets.QCheckBox("Use video adaptation")
self.adapt_checkbox.setChecked(True)

pseudo_threshold_label = QtWidgets.QLabel("Pseudo-label confidence threshold")
self.pseudo_threshold_label = QtWidgets.QLabel(
"Pseudo-label confidence threshold"
)
self.pseudo_threshold_spinbox = QtWidgets.QDoubleSpinBox(
decimals=2,
minimum=0.01,
Expand All @@ -171,22 +185,24 @@ def _build_tf_attributes(self) -> None:
)
self.pseudo_threshold_spinbox.setMaximumWidth(300)

adapt_iter_label = QtWidgets.QLabel("Number of adaptation iterations")
adapt_iter_label.setMinimumWidth(300)
self.adapt_iter_label = QtWidgets.QLabel("Number of adaptation iterations")
self.adapt_iter_label.setMinimumWidth(300)
self.adapt_iter_spinbox = QtWidgets.QSpinBox()
self.adapt_iter_spinbox.setRange(100, 10000)
self.adapt_iter_spinbox.setValue(1000)
self.adapt_iter_spinbox.setSingleStep(100)
self.adapt_iter_spinbox.setGroupSeparatorShown(True)
self.adapt_iter_spinbox.setMaximumWidth(300)

self.adapt_checkbox.stateChanged.connect(self._adapt_checkbox_status_changed)

model_settings_layout.addWidget(scales_label, 1, 0)
model_settings_layout.addWidget(self.scales_line, 1, 1)
model_settings_layout.addWidget(tooltip_label, 1, 2)
model_settings_layout.addWidget(self.adapt_checkbox, 2, 0)
model_settings_layout.addWidget(pseudo_threshold_label, 3, 0)
model_settings_layout.addWidget(self.pseudo_threshold_label, 3, 0)
model_settings_layout.addWidget(self.pseudo_threshold_spinbox, 3, 1)
model_settings_layout.addWidget(adapt_iter_label, 4, 0)
model_settings_layout.addWidget(self.adapt_iter_label, 4, 0)
model_settings_layout.addWidget(self.adapt_iter_spinbox, 4, 1)
self.tf_widget = QtWidgets.QWidget()
self.tf_widget.setLayout(model_settings_layout)
Expand All @@ -199,8 +215,10 @@ def _build_torch_attributes(self) -> None:
self.torch_adapt_checkbox = QtWidgets.QCheckBox("Use video adaptation")
self.torch_adapt_checkbox.setChecked(True)

pseudo_threshold_label = QtWidgets.QLabel("Pseudo-label confidence threshold")
pseudo_threshold_label.setMinimumWidth(300)
self.torch_pseudo_threshold_label = QtWidgets.QLabel(
"Pseudo-label confidence threshold"
)
self.torch_pseudo_threshold_label.setMinimumWidth(300)
self.torch_pseudo_threshold_spinbox = QtWidgets.QDoubleSpinBox(
decimals=2,
minimum=0.01,
Expand All @@ -211,32 +229,74 @@ def _build_torch_attributes(self) -> None:
)
self.torch_pseudo_threshold_spinbox.setMaximumWidth(300)

adapt_epoch_label = QtWidgets.QLabel("Number of adaptation epochs")
adapt_epoch_label.setMinimumWidth(300)
self.torch_adapt_epoch_label = QtWidgets.QLabel("Number of adaptation epochs")
self.torch_adapt_epoch_label.setMinimumWidth(300)
self.torch_adapt_epoch_spinbox = QtWidgets.QSpinBox()
self.torch_adapt_epoch_spinbox.setRange(1, 50)
self.torch_adapt_epoch_spinbox.setValue(4)
self.torch_adapt_epoch_spinbox.setMaximumWidth(300)

adapt_det_epoch_label = QtWidgets.QLabel("Number of detector adaptation epochs")
adapt_det_epoch_label.setMinimumWidth(300)
self.torch_adapt_det_epoch_label = QtWidgets.QLabel(
"Number of detector adaptation epochs"
)
self.torch_adapt_det_epoch_label.setMinimumWidth(300)
self.torch_adapt_det_epoch_spinbox = QtWidgets.QSpinBox()
self.torch_adapt_det_epoch_spinbox.setRange(1, 50)
self.torch_adapt_det_epoch_spinbox.setValue(4)
self.torch_adapt_det_epoch_spinbox.setMaximumWidth(300)

self.torch_adapt_checkbox.stateChanged.connect(
self._torch_adapt_checkbox_status_changed
)

self.detector_batch_size_combo_label = QtWidgets.QLabel("Detector batch size")
self.detector_batch_size_combo = QtWidgets.QComboBox()
self.detector_batch_size_combo.setMinimumWidth(100)
self.detector_batch_size_combo.addItems([str(2**i) for i in range(6)])
self.detector_batch_size_combo.setCurrentIndex(0)

torch_settings_layout.addWidget(self.torch_adapt_checkbox, 1, 0)
torch_settings_layout.addWidget(pseudo_threshold_label, 2, 0)
torch_settings_layout.addWidget(self.torch_pseudo_threshold_label, 2, 0)
torch_settings_layout.addWidget(self.torch_pseudo_threshold_spinbox, 2, 1)
torch_settings_layout.addWidget(adapt_epoch_label, 3, 0)
torch_settings_layout.addWidget(self.torch_adapt_epoch_label, 3, 0)
torch_settings_layout.addWidget(self.torch_adapt_epoch_spinbox, 3, 1)
torch_settings_layout.addWidget(adapt_det_epoch_label, 4, 0)
torch_settings_layout.addWidget(self.torch_adapt_det_epoch_label, 4, 0)
torch_settings_layout.addWidget(self.torch_adapt_det_epoch_spinbox, 4, 1)
torch_settings_layout.addWidget(self.detector_batch_size_combo_label, 5, 0)
torch_settings_layout.addWidget(self.detector_batch_size_combo, 5, 1)
self.torch_widget = QtWidgets.QWidget()
self.torch_widget.setLayout(torch_settings_layout)
self.torch_widget.hide()
self.main_layout.addWidget(self.torch_widget)

def _adapt_checkbox_status_changed(self, state: int) -> None:
if Qt.CheckState(state) == Qt.Checked:
self.pseudo_threshold_label.show()
self.pseudo_threshold_spinbox.show()
self.adapt_iter_label.show()
self.adapt_iter_spinbox.show()
else:
self.pseudo_threshold_label.hide()
self.pseudo_threshold_spinbox.hide()
self.adapt_iter_label.hide()
self.adapt_iter_spinbox.hide()

def _torch_adapt_checkbox_status_changed(self, state: int) -> None:
if Qt.CheckState(state) == Qt.Checked:
self.torch_pseudo_threshold_label.show()
self.torch_pseudo_threshold_spinbox.show()
self.torch_adapt_epoch_label.show()
self.torch_adapt_epoch_spinbox.show()
self.torch_adapt_det_epoch_label.show()
self.torch_adapt_det_epoch_spinbox.show()
else:
self.torch_pseudo_threshold_label.hide()
self.torch_pseudo_threshold_spinbox.hide()
self.torch_adapt_epoch_label.hide()
self.torch_adapt_epoch_spinbox.hide()
self.torch_adapt_det_epoch_label.hide()
self.torch_adapt_det_epoch_spinbox.hide()

def select_folder(self):
dirname = QtWidgets.QFileDialog.getExistingDirectory(
self, "Please select a folder", self.root.project_folder
Expand Down Expand Up @@ -270,7 +330,7 @@ def _handle_validation_change(self, state):
self.scales_line.setStyleSheet(f"border: 1px solid {color}")
QTimer.singleShot(500, lambda: self.scales_line.setStyleSheet(""))

def run_video_adaptation(self):
def run_video_inference_superanimal(self):
videos = list(self.files)
if not videos:
msg = QtWidgets.QMessageBox()
Expand All @@ -284,6 +344,9 @@ def run_video_adaptation(self):

supermodel_name = self.model_combo.currentText()
videotype = self.video_selection_widget.videotype_widget.currentText()
create_labeled_video = self.create_labeled_video_checkbox.isChecked()
batch_size = int(self.batch_size_combo.currentText())
detector_batch_size = int(self.detector_batch_size_combo.currentText())
kwargs = self._gather_kwargs()

can_run_in_background = False
Expand All @@ -294,6 +357,9 @@ def run_video_adaptation(self):
supermodel_name,
videotype=videotype,
dest_folder=self._destfolder,
create_labeled_video=create_labeled_video,
batch_size=batch_size,
detector_batch_size=detector_batch_size,
**kwargs,
)

Expand All @@ -309,6 +375,9 @@ def run_video_adaptation(self):
supermodel_name,
videotype=videotype,
dest_folder=self._destfolder,
create_labeled_video=create_labeled_video,
batch_size=batch_size,
detector_batch_size=detector_batch_size,
**kwargs,
)
self.signal_analysis_complete()
Expand Down Expand Up @@ -347,56 +416,63 @@ def _gather_kwargs(self) -> dict:
def _update_available_models(self, engine: Engine) -> None:
current_dataset = self.model_combo.currentText()

while self.model_combo.count() > 0:
self.model_combo.removeItem(0)

if engine == Engine.TF:
supermodels = ["superanimal_topviewmouse", "superanimal_quadruped"]
else:
supermodels = dlclibrary.get_available_datasets()

self.model_combo.addItems(supermodels)
if current_dataset in supermodels:
self.model_combo.setCurrentIndex(supermodels.index(current_dataset))
set_combo_items(
combo_box=self.model_combo,
items=supermodels,
index=(
supermodels.index(current_dataset)
if current_dataset in supermodels
else 0
),
)

def _update_pose_models(self, super_animal: str) -> None:
while self.net_type_selector.count() > 0:
self.net_type_selector.removeItem(0)

if len(super_animal) == 0:
set_combo_items(combo_box=self.net_type_selector, items=[])
return

if self.root.engine == Engine.TF:
self.net_type_selector.addItems(["dlcrnet"])
else:
self.net_type_selector.addItems(
dlclibrary.get_available_models(super_animal)
)
set_combo_items(
combo_box=self.net_type_selector,
items=(
["dlcrnet"]
if self.root.engine == Engine.TF
else dlclibrary.get_available_models(super_animal)
),
)

def _update_detectors(self, super_animal: str) -> None:
while self.detector_type_selector.count() > 0:
self.detector_type_selector.removeItem(0)

if len(super_animal) == 0:
set_combo_items(combo_box=self.detector_type_selector, items=[])
return

if self.root.engine == Engine.TF:
self.detector_type_selector.addItems(["dlcrnet"])
else:
self.detector_type_selector.addItems(
dlclibrary.get_available_detectors(super_animal)
)
set_combo_items(
combo_box=self.detector_type_selector,
items=(
[]
if self.root.engine == Engine.TF
else dlclibrary.get_available_detectors(super_animal)
),
)

@Slot(Engine)
def _on_engine_change(self, engine: Engine) -> None:
self._update_available_models(engine)
if engine == Engine.PYTORCH:
self.tf_widget.hide()
self.torch_widget.show()
self.detector_type_text.show()
self.detector_type_selector.show()
self.torch_widget.show()
self.detector_batch_size_combo_label.show()
self.detector_batch_size_combo.show()
else:
self.tf_widget.show()
self.torch_widget.hide()
self.detector_type_text.hide()
self.detector_type_selector.hide()
self.tf_widget.show()
self.detector_batch_size_combo_label.hide()
self.detector_batch_size_combo.hide()
7 changes: 7 additions & 0 deletions deeplabcut/modelzoo/video_inference.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,7 @@ def video_inference_superanimal(
customized_detector_checkpoint: Optional[str] = None,
customized_model_config: Optional[str] = None,
plot_bboxes: bool = True,
create_labeled_video: bool = True,
):
"""
This function performs inference on videos using a pretrained SuperAnimal model.
Expand Down Expand Up @@ -172,6 +173,9 @@ def video_inference_superanimal(
plot_bboxes (bool):
If using Top-Down approach, whether to plot the detector's bounding boxes. The default is True.

create_labeled_video (bool):
Specifies if a labeled video needs to be created, True by default.

Raises:
NotImplementedError:
If the model is not found in the modelzoo.
Expand Down Expand Up @@ -320,6 +324,7 @@ def video_inference_superanimal(
pcutoff,
adapt_iterations,
pseudo_threshold,
create_labeled_video=create_labeled_video,
)
elif framework == "pytorch":
if detector_name is None:
Expand Down Expand Up @@ -383,6 +388,7 @@ def video_inference_superanimal(
output_suffix=output_suffix,
plot_bboxes=plot_bboxes,
bboxes_pcutoff=bbox_threshold,
create_labeled_video=create_labeled_video,
)

# we prepare the pseudo dataset in the same folder of the target video
Expand Down Expand Up @@ -558,4 +564,5 @@ def video_inference_superanimal(
output_suffix=output_suffix,
plot_bboxes=plot_bboxes,
bboxes_pcutoff=bbox_threshold,
create_labeled_video=create_labeled_video,
)
30 changes: 17 additions & 13 deletions deeplabcut/pose_estimation_pytorch/modelzoo/inference.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ def _video_inference_superanimal(
output_suffix: str = "",
plot_bboxes: bool = True,
bboxes_pcutoff: float = 0.9,
create_labeled_video: bool = True,
) -> dict:
"""
Perform inference on a video using a superanimal model from the model zoo specified by `superanimal_name`.
Expand Down Expand Up @@ -89,6 +90,8 @@ def _video_inference_superanimal(
dest_folder: Destination folder for the results. If not specified, the
results are saved in the same folder as the video. Defaults to None.
output_suffix: The suffix to add to output file names (e.g. _before_adapt)
create_labeled_video (bool):
Specifies if a labeled video needs to be created, True by default.

Returns:
results: Dictionary with the result pd.DataFrame for each video
Expand Down Expand Up @@ -172,18 +175,19 @@ def _video_inference_superanimal(
superanimal_colormaps = get_superanimal_colormaps()
colormap = superanimal_colormaps[superanimal_name]

create_video(
video_path,
output_h5,
pcutoff=pcutoff,
fps=video.fps,
bbox=bbox,
cmap=colormap,
output_path=str(output_video),
plot_bboxes=plot_bboxes,
bboxes_list=bboxes_list,
bboxes_pcutoff=bboxes_pcutoff,
)
print(f"Video with predictions was saved as {output_path}")
if create_labeled_video:
create_video(
video_path,
output_h5,
pcutoff=pcutoff,
fps=video.fps,
bbox=bbox,
cmap=colormap,
output_path=str(output_video),
plot_bboxes=plot_bboxes,
bboxes_list=bboxes_list,
bboxes_pcutoff=bboxes_pcutoff,
)
print(f"Video with predictions was saved as {output_path}")

return results
Loading