DeepLabCut · AlexEMG · Mar 7, 2023 · Oct 14, 2022 · Oct 15, 2022 · Oct 20, 2022
diff --git a/deeplabcut/__init__.py b/deeplabcut/__init__.py
@@ -107,6 +107,7 @@
     visualize_paf,
     extract_save_all_maps,
     export_model,
+    video_inference_superanimal
 )
 
 

diff --git a/deeplabcut/gui/tabs/__init__.py b/deeplabcut/gui/tabs/__init__.py
@@ -17,6 +17,7 @@
 from deeplabcut.gui.tabs.extract_outlier_frames import ExtractOutlierFrames
 from deeplabcut.gui.tabs.label_frames import LabelFrames
 from deeplabcut.gui.tabs.manage_project import ManageProject
+from deeplabcut.gui.tabs.modelzoo import ModelZoo
 from deeplabcut.gui.tabs.open_project import OpenProject
 from deeplabcut.gui.tabs.refine_tracklets import RefineTracklets
 from deeplabcut.gui.tabs.train_network import TrainNetwork

diff --git a/deeplabcut/gui/tabs/modelzoo.py b/deeplabcut/gui/tabs/modelzoo.py
@@ -0,0 +1,103 @@
+import deeplabcut
+from PySide6 import QtWidgets
+from PySide6.QtCore import Qt, Signal, QTimer
+from PySide6.QtGui import QRegularExpressionValidator
+from deeplabcut.gui.components import (
+    DefaultTab,
+    VideoSelectionWidget,
+    _create_label_widget,
+    _create_grid_layout,
+)
+from deeplabcut.modelzoo.utils import parse_available_supermodels
+
+
+class RegExpValidator(QRegularExpressionValidator):
+    validationChanged = Signal(QRegularExpressionValidator.State)
+
+    def validate(self, input_, pos):
+        state, input_, pos = super().validate(input_, pos)
+        self.validationChanged.emit(state)
+        return state, input_, pos
+
+
+class ModelZoo(DefaultTab):
+    def __init__(self, root, parent, h1_description):
+        super().__init__(root, parent, h1_description)
+        self._val_pattern = "(\d{3,5},\s*)+\d{3,5}"
+        self._set_page()
+
+    @property
+    def files(self):
+        return self.video_selection_widget.files
+
+    def _set_page(self):
+        self.main_layout.addWidget(_create_label_widget("Video Selection", "font:bold"))
+        self.video_selection_widget = VideoSelectionWidget(self.root, self)
+        self.main_layout.addWidget(self.video_selection_widget)
+
+        model_settings_layout = _create_grid_layout(margins=(20, 0, 0, 0))
+
+        section_title = _create_label_widget(
+            "Supermodel Settings", "font:bold", (0, 50, 0, 0)
+        )
+
+        model_combo_text = QtWidgets.QLabel("Supermodel name")
+        self.model_combo = QtWidgets.QComboBox()
+        supermodels = parse_available_supermodels()
+        self.model_combo.addItems(supermodels.keys())
+
+        scales_label = QtWidgets.QLabel('Scale list')
+        self.scales_line = QtWidgets.QLineEdit("", parent=self)
+        self.scales_line.setPlaceholderText("Optionally input a list of integer sizes separated by commas...")
+        validator = RegExpValidator(self._val_pattern, self)
+        validator.validationChanged.connect(self._handle_validation_change)
+        self.scales_line.setValidator(validator)
+
+        model_settings_layout.addWidget(section_title, 0, 0)
+        model_settings_layout.addWidget(model_combo_text, 1, 0)
+        model_settings_layout.addWidget(self.model_combo, 1, 1)
+        model_settings_layout.addWidget(scales_label, 2, 0)
+        model_settings_layout.addWidget(self.scales_line, 2, 1)
+        self.main_layout.addLayout(model_settings_layout)
+
+        self.run_button = QtWidgets.QPushButton("Run")
+        self.run_button.clicked.connect(self.run_video_adaptation)
+        self.main_layout.addWidget(self.run_button, alignment=Qt.AlignRight)
+
+    def _handle_validation_change(self, state):
+        if state == RegExpValidator.Invalid:
+            color = 'red'
+        elif state == RegExpValidator.Intermediate:
+            color = 'gold'
+        elif state == RegExpValidator.Acceptable:
+            color = 'lime'
+        self.scales_line.setStyleSheet(f'border: 1px solid {color}')
+        QTimer.singleShot(500, lambda: self.scales_line.setStyleSheet(''))
+
+    def run_video_adaptation(self):
+        videos = list(self.files)
+        if not videos:
+            msg = QtWidgets.QMessageBox()
+            msg.setIcon(QtWidgets.QMessageBox.Critical)
+            msg.setText("You must select a video file")
+            msg.setWindowTitle("Error")
+            msg.setMinimumWidth(400)
+            msg.setStandardButtons(QtWidgets.QMessageBox.Ok)
+            msg.exec_()
+            return
+
+        scales = []
+        scales_ = self.scales_line.text()
+        if scales_:
+            if (self.scales_line.validator().validate(scales_, 0)[0] == RegExpValidator.Acceptable):
+                scales = list(map(int, scales_.split(",")))
+        supermodel_name = self.model_combo.currentText()
+        videotype = self.video_selection_widget.videotype_widget.currentText()
+
+        deeplabcut.video_inference_superanimal(
+            videos,
+            supermodel_name,
+            videotype=videotype,
+            video_adapt=True,
+            scale_list=scales,
+        )
diff --git a/deeplabcut/gui/window.py b/deeplabcut/gui/window.py
@@ -471,6 +471,9 @@ def add_tabs(self):
         self.refine_tracklets = RefineTracklets(
             root=self, parent=None, h1_description="DeepLabCut - Refine labels"
         )
+        self.modelzoo = ModelZoo(
+            root=self, parent=None, h1_description="DeepLabCut - Model Zoo"
+        )
         self.video_editor = VideoEditor(
             root=self, parent=None, h1_description="DeepLabCut - Optional Video Editor"
         )
@@ -490,6 +493,7 @@ def add_tabs(self):
             self.extract_outlier_frames, "Extract outlier frames (*)"
         )
         self.tab_widget.addTab(self.refine_tracklets, "Refine tracklets (*)")
+        self.tab_widget.addTab(self.modelzoo, "Model Zoo")
         self.tab_widget.addTab(self.video_editor, "Video editor (*)")
 
         if not self.is_multianimal:

diff --git a/deeplabcut/modelzoo/__init__.py b/deeplabcut/modelzoo/__init__.py
diff --git a/deeplabcut/modelzoo/api/__init__.py b/deeplabcut/modelzoo/api/__init__.py
@@ -0,0 +1 @@
+from .spatiotemporal_adapt import SpatiotemporalAdaptation
diff --git a/deeplabcut/modelzoo/api/spatiotemporal_adapt.py b/deeplabcut/modelzoo/api/spatiotemporal_adapt.py
@@ -0,0 +1,199 @@
+import deeplabcut
+import glob
+import os
+from deeplabcut.modelzoo.utils import parse_available_supermodels
+from deeplabcut.modelzoo.api import superanimal_inference
+from deeplabcut.utils.plotting import _plot_trajectories
+from pathlib import Path
+
+
+class SpatiotemporalAdaptation:
+    def __init__(
+        self,
+        video_path,
+        supermodel_name,
+        scale_list=[],
+        videotype="mp4",
+        adapt_iterations=1000,
+        modelfolder="",
+        customized_pose_config="",
+        init_weights="",
+    ):
+
+        """
+        This class supports video adaptation to a super model.
+
+        Parameters
+        ----------
+        video_path: string
+           The string to the path of the video
+        init_weights: string
+           The path to a superanimal model's checkpoint
+        supermodel_name: string
+           Currently we support supertopview(LabMice) and superquadruped (quadruped side-view animals)
+        scale_list: list
+           A list of different resolutions for the spatial pyramid
+        videotype: string
+           Checks for the extension of the video in case the input to the video is a directory.\n Only videos with this extension are analyzed. The default is ``.avi``
+        adapt_iterations: int
+           Number of iterations for adaptation training. Empirically 1000 is sufficient. Training longer can cause worse performance depending whether there is occlusion in the video
+        modelfolder: string, optional
+           Because the API does not need a dlc project, the checkpoint and logs go to this temporary model folder, and otherwise model is saved to the current work place
+        customized_pose_config: string, optional
+           For future support of non modelzoo model
+
+        Examples
+        --------
+
+        from  deeplabcut.modelzoo.apis import SpatiotemporalAdaptation
+        video_path = '/mnt/md0/shaokai/openfield_video/m3v1mp4.mp4'
+        superanimal_name = 'superanimal_topviewmouse'
+        videotype = 'mp4'
+        >>> adapter = SpatiotemporalAdaptation(video_path,
+                                       superanimal_name,
+                                       modelfolder = "temp_topview",
+                                       videotype = videotype)
+
+        adapter.before_adapt_inference()
+        adapter.adaptation_training()
+        adapter.after_adapt_inference()
+
+
+        """
+        supermodels = parse_available_supermodels()
+        if supermodel_name not in supermodels:
+            raise ValueError(f"`supermodel_name` should be one of: {', '.join(supermodels)}.")
+
+        self.video_path = video_path
+        self.supermodel_name = supermodel_name
+        self.scale_list = scale_list
+        self.videotype = videotype
+        vname = str(Path(self.video_path).stem)
+        self.adapt_modelprefix = vname + "_video_adaptation"
+        self.adapt_iterations = adapt_iterations
+        self.modelfolder = modelfolder
+        self.init_weights = init_weights
+
+        if not customized_pose_config:
+            dlc_root_path = os.sep.join(deeplabcut.__file__.split(os.sep)[:-1])
+            self.customized_pose_config = os.path.join(
+                dlc_root_path,
+                "pose_estimation_tensorflow",
+                "superanimal_configs",
+                supermodels[self.supermodel_name],
+            )
+        else:
+            self.customized_pose_config = customized_pose_config
+
+    def before_adapt_inference(self,
+                               make_video=False,
+                               **kwargs):
+        if self.init_weights != "":
+            _ = superanimal_inference.video_inference(
+                [self.video_path],
+                self.supermodel_name,
+                videotype=self.videotype,
+                scale_list=self.scale_list,
+                init_weights=self.init_weights,
+                customized_test_config=self.customized_pose_config,
+            )
+        else:
+            self.init_weights, _ = superanimal_inference.video_inference(
+                [self.video_path],
+                self.supermodel_name,
+                videotype=self.videotype,
+                scale_list=self.scale_list,
+                customized_test_config=self.customized_pose_config,
+            )
+        if make_video:
+            deeplabcut.create_labeled_video(
+                "",
+                [self.video_path],
+                videotype=self.videotype,
+                filtered=False,
+                init_weights=self.init_weights,
+                draw_skeleton=True,
+                superanimal_name=self.supermodel_name,
+                **kwargs
+            )
+
+    def train_without_project(self, pseudo_label_path, **kwargs):
+        from deeplabcut.pose_estimation_tensorflow.core.train_multianimal import train
+
+        displayiters = kwargs.pop("displayiters", 500)
+        saveiters = kwargs.pop("saveiters", 1000)
+        train(
+            self.customized_pose_config,
+            displayiters=displayiters,
+            saveiters=saveiters,
+            maxiters=self.adapt_iterations,
+            modelfolder=self.modelfolder,
+            init_weights=self.init_weights,
+            pseudo_labels=pseudo_label_path,
+            video_path=self.video_path,
+            **kwargs
+        )
+
+    def adaptation_training(self, displayiters=500, saveiters=1000, **kwargs):
+        """
+        There should be two choices, either taking a config, with is then assuming there is a DLC project.
+        Or we make up a fake one, then we use a light way convention to do adaptation
+        """
+
+        # looking for the pseudo label path
+        DLCscorer = "DLC_" + Path(self.init_weights).stem
+        vname = str(Path(self.video_path).stem)
+        video_root = Path(self.video_path).parent
+
+        _, pseudo_label_path, _, _ = deeplabcut.auxiliaryfunctions.load_analyzed_data(
+            video_root, vname, DLCscorer, False, ""
+        )
+        if self.modelfolder != "":
+            os.makedirs(self.modelfolder, exist_ok=True)
+
+        self.train_without_project(
+            pseudo_label_path,
+            displayiters=displayiters,
+            saveiters=saveiters,
+            **kwargs,
+        )
+
+    def after_adapt_inference(self, **kwargs):
+
+        pattern = os.path.join(
+            self.modelfolder, f"snapshot-{self.adapt_iterations}.index"
+        )
+        ref_proj_config_path = ""
+
+        files = glob.glob(pattern)
+
+        if not len(files):
+            raise ValueError("Weights were not found.")
+
+        adapt_weights = files[0].replace(".index", "")
+
+        # spatial pyramid is not for adapted model
+
+        scale_list = kwargs.pop('scale_list', [])
+        _, datafiles = superanimal_inference.video_inference(
+            [self.video_path],
+            self.supermodel_name,
+            videotype=self.videotype,
+            init_weights=adapt_weights,
+            scale_list=scale_list,
+            customized_test_config=self.customized_pose_config,
+        )
+
+        if kwargs.pop('plot_trajectories', True):
+            _plot_trajectories(datafiles[0])
+
+        deeplabcut.create_labeled_video(
+            ref_proj_config_path,
+            [self.video_path],
+            videotype=self.videotype,
+            filtered=False,
+            init_weights=adapt_weights,
+            draw_skeleton=True,
+            superanimal_name=self.supermodel_name,
+            **kwargs
+        )
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		from .spatiotemporal_adapt import SpatiotemporalAdaptation