dialog

sleap.gui.learning.dialog

Dialogs for running training and/or inference in GUI.

Classes:

Name	Description
LearningDialog	Dialog for running training and/or inference.
TrainingEditorWidget	Dialog for viewing and modifying training profiles (model hyperparameters).
TrainingPipelineWidget	Widget used in :py:class:LearningDialog for configuring pipeline.

LearningDialog

Bases: QDialog

Dialog for running training and/or inference.

The dialog shows tabs for configuring the pipeline ( :py:class:TrainingPipelineWidget) and, depending on the pipeline, for each specific model (:py:class:TrainingEditorWidget).

In training mode, the model hyperpameters are editable unless you're using a trained model; they are read-only in inference mode.

Parameters:

Name	Type	Description	Default
mode	Text	either "training" or "inference".	required
labels_filename	Text	path to labels file, used for default location to save models.	required
labels	Optional[Labels]	the Labels object (can also be loaded from given filename)	None
skeleton	Optional[Skeleton]	the Skeleton object (can also be taken from Labels), used for list of nodes for (e.g.) selecting anchor node	None

Methods:

Name	Description
copy	Copy scripts and configs to clipboard
export_package	Export training job package.
run	Run with current dialog settings.
save	Save scripts and configs to run pipeline.
update_loaded_config	Update a loaded preset config with values from the training editor.

Attributes:

Name	Type	Description
frame_selection	Dict[str, Dict[Video, List[int]]]	Returns dictionary with frames that user has selected for learning.

Source code in sleap/gui/learning/dialog.py

class LearningDialog(QtWidgets.QDialog):
    """
    Dialog for running training and/or inference.

    The dialog shows tabs for configuring the pipeline (
    :py:class:`TrainingPipelineWidget`) and, depending on the pipeline, for
    each specific model (:py:class:`TrainingEditorWidget`).

    In training mode, the model hyperpameters are editable unless you're using
    a trained model; they are read-only in inference mode.

    Arguments:
        mode: either "training" or "inference".
        labels_filename: path to labels file, used for default location to
            save models.
        labels: the `Labels` object (can also be loaded from given filename)
        skeleton: the `Skeleton` object (can also be taken from `Labels`), used
            for list of nodes for (e.g.) selecting anchor node
    """

    _handle_learning_finished = QtCore.Signal(int)

    def __init__(
        self,
        mode: Text,
        labels_filename: Text,
        labels: Optional[Labels] = None,
        skeleton: Optional["Skeleton"] = None,
        *args,
        **kwargs,
    ):
        super(LearningDialog, self).__init__()

        if labels is None:
            labels = load_file(labels_filename)

        if skeleton is None and labels.skeletons:
            skeleton = labels.skeletons[0]

        self.mode = mode
        self.labels_filename = labels_filename
        self.labels = labels
        self.skeleton = skeleton

        self._frame_selection = None

        self.current_pipeline = ""

        self.tabs: Dict[str, TrainingEditorWidget] = dict()
        self.shown_tab_names = []

        self._cfg_getter = configs.TrainingConfigsGetter.make_from_labels_filename(
            labels_filename=self.labels_filename
        )

        # Layout for buttons
        buttons = QtWidgets.QDialogButtonBox()
        self.copy_button = buttons.addButton(
            "Copy to clipboard", QtWidgets.QDialogButtonBox.ActionRole
        )
        self.save_button = buttons.addButton(
            "Save configuration files...", QtWidgets.QDialogButtonBox.ActionRole
        )
        self.export_button = buttons.addButton(
            "Export training job package...", QtWidgets.QDialogButtonBox.ActionRole
        )
        self.cancel_button = buttons.addButton(QtWidgets.QDialogButtonBox.Cancel)
        self.run_button = buttons.addButton("Run", QtWidgets.QDialogButtonBox.ApplyRole)

        self.copy_button.setToolTip("Copy configuration to the clipboard")
        self.save_button.setToolTip("Save scripts and configuration to run pipeline.")
        self.export_button.setToolTip(
            "Export data, configuration, and scripts for remote training and inference."
        )
        self.run_button.setToolTip("Run pipeline locally (GPU recommended).")

        buttons_layout = QtWidgets.QHBoxLayout()
        buttons_layout.addWidget(buttons, alignment=QtCore.Qt.AlignTop)

        buttons_layout_widget = QtWidgets.QWidget()
        buttons_layout_widget.setLayout(buttons_layout)

        self.pipeline_form_widget = TrainingPipelineWidget(mode=mode, skeleton=skeleton)
        if mode == "training":
            tab_label = "Training Pipeline"
        elif mode == "inference":
            # self.pipeline_form_widget = InferencePipelineWidget()
            tab_label = "Inference Pipeline"
        else:
            raise ValueError(f"Invalid LearningDialog mode: {mode}")

        self.tab_widget = QtWidgets.QTabWidget()

        self.tab_widget.addTab(self.pipeline_form_widget, tab_label)
        self.make_tabs()

        self.message_widget = QtWidgets.QLabel("")

        # Layout for entire dialog
        content_widget = QtWidgets.QWidget()
        content_layout = QtWidgets.QVBoxLayout(content_widget)

        content_layout.addWidget(self.tab_widget)
        content_layout.addWidget(self.message_widget)
        content_layout.addWidget(buttons_layout_widget)

        # Create the QScrollArea.
        scroll_area = QtWidgets.QScrollArea()
        scroll_area.setWidgetResizable(True)
        scroll_area.setWidget(content_widget)

        scroll_area.setVerticalScrollBarPolicy(QtCore.Qt.ScrollBarAsNeeded)
        scroll_area.setHorizontalScrollBarPolicy(QtCore.Qt.ScrollBarAsNeeded)

        layout = QtWidgets.QVBoxLayout(self)
        layout.addWidget(scroll_area)

        self.adjust_initial_size()

        # Default to most recently trained pipeline (if there is one)
        self.set_default_pipeline_tab()

        # Connect functions to update pipeline tabs when pipeline changes
        self.pipeline_form_widget.updatePipeline.connect(self.set_pipeline)
        self.pipeline_form_widget.emitPipeline()

        self.connect_signals()

        # Connect actions for buttons
        self.copy_button.clicked.connect(self.copy)
        self.save_button.clicked.connect(lambda: self.save())
        self.export_button.clicked.connect(lambda: self.export_package())
        self.cancel_button.clicked.connect(self.reject)
        self.run_button.clicked.connect(self.run)

    def adjust_initial_size(self):
        # Get screen size
        screen = QtGui.QGuiApplication.primaryScreen().availableGeometry()

        max_width = 1860
        max_height = 1150
        margin = 0.10

        # Calculate target width and height
        target_width = min(screen.width() - screen.width() * margin, max_width)
        target_height = min(screen.height() - screen.height() * margin, max_height)
        # Set the dialog's dimensions
        self.resize(target_width, target_height)

    def update_file_lists(self):
        self._cfg_getter.update()
        for tab in self.tabs.values():
            tab.update_file_list()

    @staticmethod
    def count_total_frames_for_selection_option(
        videos_frames: Dict[Video, List[int]],
    ) -> int:
        if not videos_frames:
            return 0

        count = 0
        for frame_list in videos_frames.values():
            # Check for [X, Y) range given as (X, -Y) tuple
            if len(frame_list) == 2 and frame_list[1] < 0:
                count += -frame_list[1] - frame_list[0]
            elif frame_list != (0, 0):
                count += len(frame_list)

        return count

    @property
    def frame_selection(self) -> Dict[str, Dict[Video, List[int]]]:
        """
        Returns dictionary with frames that user has selected for learning.
        """
        return self._frame_selection

    @frame_selection.setter
    def frame_selection(self, frame_selection: Dict[str, Dict[Video, List[int]]]):
        """Sets options of frames on which to run learning."""
        self._frame_selection = frame_selection

        if "_predict_frames" in self.pipeline_form_widget.fields.keys():
            prediction_options = []

            total_random = 0
            total_suggestions = 0
            total_user = 0
            random_video = 0
            clip_length = 0
            video_length = 0
            all_videos_length = 0

            # Determine which options are available given _frame_selection
            if "random" in self._frame_selection:
                total_random = self.count_total_frames_for_selection_option(
                    self._frame_selection["random"]
                )
            if "random_video" in self._frame_selection:
                random_video = self.count_total_frames_for_selection_option(
                    self._frame_selection["random_video"]
                )
            if "suggestions" in self._frame_selection:
                total_suggestions = self.count_total_frames_for_selection_option(
                    self._frame_selection["suggestions"]
                )
            if "user" in self._frame_selection:
                total_user = self.count_total_frames_for_selection_option(
                    self._frame_selection["user"]
                )
            if "clip" in self._frame_selection:
                clip_length = self.count_total_frames_for_selection_option(
                    self._frame_selection["clip"]
                )
            if "video" in self._frame_selection:
                video_length = self.count_total_frames_for_selection_option(
                    self._frame_selection["video"]
                )
            if "all_videos" in self._frame_selection:
                all_videos_length = self.count_total_frames_for_selection_option(
                    self._frame_selection["all_videos"]
                )

            # Build list of options
            # Priority for default (lowest to highest):
            #   "nothing" (if training)
            #   "current frame" (if inference)
            #   "suggested frames" (if available)
            #   "selected clip" (if available)
            if self.mode != "inference":
                prediction_options.append("nothing")
            prediction_options.append("current frame")
            default_option = "nothing" if self.mode != "inference" else "current frame"

            option = f"random frames ({total_random} total frames)"
            prediction_options.append(option)

            if random_video > 0:
                option = f"random frames in current video ({random_video} frames)"
                prediction_options.append(option)

            if total_suggestions > 0:
                option = f"suggested frames ({total_suggestions} total frames)"
                prediction_options.append(option)
                default_option = option

            if total_user > 0:
                option = f"user labeled frames ({total_user} total frames)"
                prediction_options.append(option)

            if clip_length > 0:
                option = f"selected clip ({clip_length} frames)"
                prediction_options.append(option)
                default_option = option

            prediction_options.append(f"entire current video ({video_length} frames)")

            if len(self.labels.videos) > 1:
                prediction_options.append(f"all videos ({all_videos_length} frames)")

            self.pipeline_form_widget.fields["_predict_frames"].set_options(
                prediction_options, default_option
            )

    def connect_signals(self):
        self.pipeline_form_widget.valueChanged.connect(self.on_tab_data_change)

        for head_name, tab in self.tabs.items():
            tab.valueChanged.connect(lambda n=head_name: self.on_tab_data_change(n))

    def disconnect_signals(self):
        self.pipeline_form_widget.valueChanged.disconnect()

        for head_name, tab in self.tabs.items():
            tab.valueChanged.disconnect()

    def make_tabs(self):
        heads = (
            "single_instance",
            "centroid",
            "centered_instance",
            "bottomup",
            "multi_class_topdown",
            "multi_class_bottomup",
        )

        video = self.labels.videos[0] if self.labels else None

        for head_name in heads:
            self.tabs[head_name] = TrainingEditorWidget(
                video=video,
                skeleton=self.skeleton,
                head=head_name,
                cfg_getter=self._cfg_getter,
                require_trained=(self.mode == "inference"),
            )

    def adjust_data_to_update_other_tabs(self, source_data, updated_data=None):
        if updated_data is None:
            updated_data = source_data

        anchor_part = None
        set_anchor = False

        if "model_config.head_configs.centroid.confmaps.anchor_part" in source_data:
            anchor_part = source_data[
                "model_config.head_configs.centroid.confmaps.anchor_part"
            ]
            set_anchor = True
        elif (
            "model_config.head_configs.centered_instance.confmaps.anchor_part"
            in source_data
        ):
            anchor_part = source_data[
                "model_config.head_configs.centered_instance.confmaps.anchor_part"
            ]
            set_anchor = True
        elif (
            "model_config.head_configs.multi_class_topdown.confmaps.anchor_part"
            in source_data
        ):
            anchor_part = source_data[
                "model_config.head_configs.multi_class_topdown.confmaps.anchor_part"
            ]
            set_anchor = True

        # Use None instead of empty string/list
        anchor_part = anchor_part or None

        if set_anchor:
            updated_data["model_config.head_configs.centroid.confmaps.anchor_part"] = (
                anchor_part
            )
            updated_data[
                "model_config.head_configs.centered_instance.confmaps.anchor_part"
            ] = anchor_part
            updated_data[
                "model_config.head_configs.multi_class_topdown.confmaps.anchor_part"
            ] = anchor_part

    def update_tabs_from_pipeline(self, source_data):
        self.adjust_data_to_update_other_tabs(source_data)

        for tab in self.tabs.values():
            tab.set_fields_from_key_val_dict(source_data)

    def update_tabs_from_tab(self, source_data):
        data_to_transfer = dict()
        self.adjust_data_to_update_other_tabs(source_data, data_to_transfer)

        if data_to_transfer:
            for tab in self.tabs.values():
                tab.set_fields_from_key_val_dict(data_to_transfer)

    def on_tab_data_change(self, tab_name=None):
        self.disconnect_signals()

        if tab_name is None:
            # Move data from pipeline tab to other tabs
            source_data = self.pipeline_form_widget.get_form_data()
            self.update_tabs_from_pipeline(source_data)
        else:
            # Get data from head-specific tab
            source_data = self.tabs[tab_name].get_all_form_data()

            self.update_tabs_from_tab(source_data)

            # Update pipeline tab
            self.pipeline_form_widget.set_form_data(source_data)

        self._validate_pipeline()

        self.connect_signals()

    def get_most_recent_pipeline_trained(self) -> Text:
        recent_cfg_info = self._cfg_getter.get_first()

        if recent_cfg_info and recent_cfg_info.head_name:
            if recent_cfg_info.head_name in ("multi_class_topdown",):
                return "top-down-id"
            if recent_cfg_info.head_name in ("centroid", "centered_instance"):
                return "top-down"
            if recent_cfg_info.head_name in ("bottomup",):
                return "bottom-up"
            if recent_cfg_info.head_name in ("single_instance",):
                return "single"
            if recent_cfg_info.head_name in ("multi_class_bottomup",):
                return "bottom-up-id"
        return ""

    def set_default_pipeline_tab(self):
        recent_pipeline_name = self.get_most_recent_pipeline_trained()
        if recent_pipeline_name:
            self.pipeline_form_widget.current_pipeline = recent_pipeline_name
        else:
            # Set default based on detection of single- vs multi-animal project.
            max_user_instance = 0
            for lf in self.labels:
                max_user_instance = max(max_user_instance, len(lf.user_instances))

            if max_user_instance == 1:
                self.pipeline_form_widget.current_pipeline = "single"
            else:
                self.pipeline_form_widget.current_pipeline = "top-down"

    def add_tab(self, tab_name):
        tab_labels = {
            "single_instance": "Single Instance Model Configuration",
            "centroid": "Centroid Model Configuration",
            "centered_instance": "Centered Instance Model Configuration",
            "bottomup": "Bottom-Up Model Configuration",
            "multi_class_topdown": "Top-Down-Id Model Configuration",
            "multi_class_bottomup": "Bottom-Up-Id Model Configuration",
        }
        self.tab_widget.addTab(self.tabs[tab_name], tab_labels[tab_name])
        self.shown_tab_names.append(tab_name)

    def remove_tabs(self):
        while self.tab_widget.count() > 1:
            self.tab_widget.removeTab(1)
        self.shown_tab_names = []

    def set_pipeline(self, pipeline: str):
        if pipeline != self.current_pipeline:
            self.remove_tabs()
            if pipeline == "top-down":
                self.add_tab("centroid")
                self.add_tab("centered_instance")
            elif pipeline == "bottom-up":
                self.add_tab("bottomup")
            elif pipeline == "top-down-id":
                self.add_tab("centroid")
                self.add_tab("multi_class_topdown")
            elif pipeline == "bottom-up-id":
                self.add_tab("multi_class_bottomup")
            elif pipeline == "single":
                self.add_tab("single_instance")
        self.current_pipeline = pipeline

        self._validate_pipeline()

    def change_tab(self, tab_idx: int):
        print(tab_idx)

    def merge_pipeline_and_head_config_data(self, head_name, head_data, pipeline_data):
        for key, val in pipeline_data.items():
            # if key.starts_with("_"):
            #     continue
            if key.startswith("model_config.head_configs."):
                key_scope = key.split(".")
                if key_scope[2] != head_name:
                    continue
            head_data[key] = val

    @staticmethod
    def update_loaded_config(
        loaded_cfg: dict, tab_cfg_key_val_dict: dict
    ):  # -> scopedkeydict.ScopedKeyDict:
        """Update a loaded preset config with values from the training editor.

        Args:
            loaded_cfg: Dict from a yaml file that was loaded from a preset or previous
                training run.
            tab_cfg_key_val_dict: A dictionary with the values extracted from the
                training editor GUI tab.

        Returns:
                    A `ScopedKeyDict` with the loaded config values overriden by the
        corresponding ones from the `tab_cfg_key_val_dict`.
        """
        # Replace params exposed in GUI with values from GUI
        for param, value in tab_cfg_key_val_dict.items():
            loaded_cfg[param] = value

        return loaded_cfg

    def get_every_head_config_data(
        self, pipeline_form_data
    ) -> List[configs.ConfigFileInfo]:
        cfg_info_list = []

        # Copy relevant data into linked fields (i.e., anchor part).
        self.adjust_data_to_update_other_tabs(pipeline_form_data)

        for tab_name in self.shown_tab_names:
            trained_cfg_info = self.tabs[tab_name].trained_config_info_to_use
            if self.tabs[tab_name].use_trained and (trained_cfg_info is not None):
                cfg_info_list.append(trained_cfg_info)

            else:
                # Get config data from GUI
                tab_cfg_key_val_dict = self.tabs[tab_name].get_all_form_data()
                self.merge_pipeline_and_head_config_data(
                    head_name=tab_name,
                    head_data=tab_cfg_key_val_dict,
                    pipeline_data=pipeline_form_data,
                )
                apply_cfg_transforms_to_key_val_dict(tab_cfg_key_val_dict)

                if trained_cfg_info is None:
                    # Config could not be loaded, just use the values from the GUI
                    loaded_cfg_scoped: dict = tab_cfg_key_val_dict
                else:
                    # Config was loaded, override with the values from the GUI
                    loaded_cfg_scoped = LearningDialog.update_loaded_config(
                        get_keyval_dict_from_omegaconf(trained_cfg_info.config),
                        tab_cfg_key_val_dict,
                    )

                # Deserialize merged dict to object
                cfg = get_omegaconf_from_gui_form(loaded_cfg_scoped)

                if len(self.labels.tracks) > 0:
                    # For multiclass topdown, the class vectors output stride
                    # should be the max stride.
                    backbone_name = find_backbone_name_from_key_val_dict(
                        tab_cfg_key_val_dict
                    )
                    max_stride = tab_cfg_key_val_dict[
                        f"model_config.backbone_config.{backbone_name}.max_stride"
                    ]

                    # Classes should be added here to prevent value error in
                    # model since we don't add them in the training config yaml.
                    if (
                        OmegaConf.select(
                            cfg,
                            "model_config.head_configs.multi_class_bottomup",
                            default=None,
                        )
                        is not None
                    ):
                        (
                            cfg.model_config.head_configs.multi_class_bottomup.class_maps.classes
                        ) = [t.name for t in self.labels.tracks]
                    elif (
                        OmegaConf.select(
                            cfg,
                            "model_config.head_configs.multi_class_topdown",
                            default=None,
                        )
                        is not None
                    ):
                        (
                            cfg.model_config.head_configs.multi_class_topdown.class_vectors.classes
                        ) = [t.name for t in self.labels.tracks]
                        (
                            cfg.model_config.head_configs.multi_class_topdown.class_vectors.output_stride
                        ) = max_stride

                cfg_info = configs.ConfigFileInfo(config=cfg, head_name=tab_name)

                cfg_info_list.append(cfg_info)

        return cfg_info_list

    def get_selected_frames_to_predict(
        self, pipeline_form_data
    ) -> Dict[Video, List[int]]:
        frames_to_predict = dict()

        if self._frame_selection is not None:
            predict_frames_choice = pipeline_form_data.get("_predict_frames", "")
            if predict_frames_choice.startswith("current frame"):
                frames_to_predict = self._frame_selection["frame"]
            elif predict_frames_choice.startswith("random frames in current video"):
                frames_to_predict = self._frame_selection["random_video"]
            elif predict_frames_choice.startswith("random"):
                frames_to_predict = self._frame_selection["random"]
            elif predict_frames_choice.startswith("selected clip"):
                frames_to_predict = self._frame_selection["clip"]
            elif predict_frames_choice.startswith("suggested"):
                frames_to_predict = self._frame_selection["suggestions"]
            elif predict_frames_choice.startswith("entire current video"):
                frames_to_predict = self._frame_selection["video"]
            elif predict_frames_choice.startswith("all videos"):
                frames_to_predict = self._frame_selection["all_videos"]
            elif predict_frames_choice.startswith("user"):
                frames_to_predict = self._frame_selection["user"]

        return frames_to_predict

    def get_items_for_inference(self, pipeline_form_data) -> runners.ItemsForInference:
        predict_frames_choice = pipeline_form_data.get("_predict_frames", "")
        batch_size = pipeline_form_data.get("batch_size")

        frame_selection = self.get_selected_frames_to_predict(pipeline_form_data)
        frame_count = self.count_total_frames_for_selection_option(frame_selection)

        if predict_frames_choice.startswith("user"):
            items_for_inference = runners.ItemsForInference(
                items=[
                    runners.DatasetItemForInference(
                        labels_path=self.labels_filename, frame_filter="user"
                    )
                ],
                total_frame_count=frame_count,
                batch_size=batch_size,
            )
        elif predict_frames_choice.startswith("suggested"):
            items_for_inference = runners.ItemsForInference(
                items=[
                    runners.DatasetItemForInference(
                        labels_path=self.labels_filename, frame_filter="suggested"
                    )
                ],
                total_frame_count=frame_count,
                batch_size=batch_size,
            )
        else:
            items_for_inference = runners.ItemsForInference.from_video_frames_dict(
                video_frames_dict=frame_selection,
                total_frame_count=frame_count,
                labels_path=self.labels_filename,
                labels=self.labels,
                batch_size=batch_size,
            )
        return items_for_inference

    def _validate_id_model(self) -> bool:
        """Make sure we have instances with tracks set for ID models."""
        if not self.labels.tracks:
            return False

        found_tracks = False
        for inst in instances(labels=self.labels):
            if type(inst) == sleap.Instance and inst.track is not None:
                found_tracks = True
                break

        return found_tracks

    def _validate_pipeline(self):
        can_run = True
        message = ""

        if self.mode == "inference":
            # Make sure we have trained models for each required head.
            untrained = [
                tab_name
                for tab_name in self.shown_tab_names
                if not self.tabs[tab_name].has_trained_config_selected
            ]
            if untrained:
                can_run = False
                message = (
                    "Cannot run inference with untrained models "
                    f"({', '.join(untrained)})."
                )
                can_run = False

        # Make sure we have instances with tracks set for ID models.
        if self.mode == "training" and self.current_pipeline in (
            "top-down-id",
            "bottom-up-id",
        ):
            can_run = self._validate_id_model()
            if not can_run:
                message = "Cannot run ID model training without tracks."

        # Make sure skeleton will be valid for bottom-up inference.
        if self.mode == "training" and self.current_pipeline == "bottom-up":
            skeleton = self.labels.skeletons[0]

            if not is_arborescence(skeleton):
                message += (
                    "Cannot run bottom-up pipeline when skeleton is not an "
                    "arborescence."
                )

                root_names = [n.name for n in root_nodes(skeleton)]
                over_max_in_degree = [n.name for n in in_degree_over_one(skeleton)]
                cycles_var = cycles(skeleton)

                if len(root_names) > 1:
                    message += (
                        f" There are multiple root nodes: {', '.join(root_names)} "
                        "(there should be exactly one node which is not a target)."
                    )

                if over_max_in_degree:
                    message += (
                        " There are nodes which are target in multiple edges: "
                        f"{', '.join(over_max_in_degree)} (maximum in-degree should be "
                        "1).</li>"
                    )

                if cycles_var:
                    cycle_strings = []
                    for cycle in cycles_var:
                        cycle_strings.append(
                            " &ndash;&gt; ".join((node.name for node in cycle))
                        )

                    message += (
                        f" There are cycles in graph: {'; '.join(cycle_strings)}."
                    )

                can_run = False

        if not can_run and message:
            message = f"<b>Unable to run:</b><br />{message}"

        self.message_widget.setText(message)
        self.run_button.setEnabled(can_run)

    def run(self):
        """Run with current dialog settings."""

        pipeline_form_data = self.pipeline_form_widget.get_form_data()

        items_for_inference = self.get_items_for_inference(pipeline_form_data)

        config_info_list = self.get_every_head_config_data(pipeline_form_data)

        # Close the dialog now that we have the data from it
        self.accept()

        # Run training/learning pipeline using the TrainingJobs
        new_counts = runners.run_learning_pipeline(
            labels_filename=self.labels_filename,
            labels=self.labels,
            config_info_list=config_info_list,
            inference_params=pipeline_form_data,
            items_for_inference=items_for_inference,
        )

        self._handle_learning_finished.emit(new_counts)

        # count < 0 means there was an error and we didn't get any results.
        if new_counts is not None and new_counts >= 0:
            total_count = items_for_inference.total_frame_count
            no_result_count = max(0, total_count - new_counts)

            message = (
                f"Inference ran on {total_count} frames."
                f"\n\nInstances were predicted on {new_counts} frames "
                f"({no_result_count} frame{'s' if no_result_count != 1 else ''} with "
                "no instances found)."
            )

            win = QtWidgets.QMessageBox(text=message)
            win.setWindowTitle("Inference Results")
            win.exec_()

    def copy(self):
        """Copy scripts and configs to clipboard"""

        # Get all info from dialog
        pipeline_form_data = self.pipeline_form_widget.get_form_data()
        config_info_list = self.get_every_head_config_data(pipeline_form_data)

        # Format information for each tab in dialog
        # output = [OmegaConf.to_yaml(pipeline_form_data)] # TODO:cfg:
        output = []
        for config_info in config_info_list:
            config_info = config_info.config
            # convert to sleap-nn cfg (yaml)
            from sleap_nn.config.training_job_config import verify_training_cfg

            config_info = filter_cfg(config_info)
            cfg = verify_training_cfg(config_info)
            cfg.data_config.train_labels_path = [self.labels_filename]
            output.append(OmegaConf.to_yaml(cfg))

        output = "\n".join(output)
        # Set the clipboard text
        clipboard = QtWidgets.QApplication.clipboard()
        clipboard.setText(output)

    def save(
        self, output_dir: Optional[str] = None, labels_filename: Optional[str] = None
    ):
        """Save scripts and configs to run pipeline."""
        if output_dir is None or not output_dir:
            labels_fn = Path(self.labels_filename)
            models_dir = Path(labels_fn.parent, "models")
            output_dir = FileDialog.openDir(
                None,
                dir=models_dir.as_posix(),
                caption="Select directory to save scripts",
            )

            if not output_dir:
                return

        pipeline_form_data = self.pipeline_form_widget.get_form_data()
        items_for_inference = self.get_items_for_inference(pipeline_form_data)
        config_info_list = self.get_every_head_config_data(pipeline_form_data)

        if labels_filename is None:
            labels_filename = self.labels_filename

        runners.write_pipeline_files(
            output_dir=output_dir,
            labels_filename=labels_filename,
            config_info_list=config_info_list,
            inference_params=pipeline_form_data,
            items_for_inference=items_for_inference,
        )

    def export_package(self, output_path: Optional[str] = None, gui: bool = True):
        """Export training job package."""
        # TODO: Warn if self.mode != "training"?
        if output_path is None or not output_path:
            # Prompt for output path.
            output_path, _ = FileDialog.save(
                caption="Export Training Job Package...",
                dir=f"{self.labels_filename}.training_job.zip",
                filter="Training Job Package (*.zip)",
            )
            if len(output_path) == 0:
                return

        # Create temp dir before packaging.
        tmp_dir = tempfile.TemporaryDirectory()

        # Remove the temp dir when program exits in case something goes wrong.
        # atexit.register(shutil.rmtree, tmp_dir.name, ignore_errors=True)

        # Check if we need to include suggestions.
        include_suggestions = False
        items_for_inference = self.get_items_for_inference(
            self.pipeline_form_widget.get_form_data()
        )
        for item in items_for_inference.items:
            if (
                isinstance(item, runners.DatasetItemForInference)
                and item.frame_filter == "suggested"
            ):
                include_suggestions = True

        # Save dataset with images.
        labels_pkg_filename = str(
            Path(self.labels_filename).with_suffix(".pkg.slp").name
        )
        if gui:
            ret = sleap.gui.commands.export_dataset_gui(
                self.labels,
                tmp_dir.name + "/" + labels_pkg_filename,
                all_labeled=False,
                suggested=include_suggestions,
            )
            if ret == "canceled":
                # Quit if user canceled during export.
                tmp_dir.cleanup()
                return
        else:
            self.labels.save(
                filename=tmp_dir.name + "/" + labels_pkg_filename,
                embed=True,
            )

        # Save config and scripts.
        self.save(tmp_dir.name, labels_filename=labels_pkg_filename)

        # Package everything.
        shutil.make_archive(
            base_name=str(Path(output_path).with_suffix("")),
            format="zip",
            root_dir=tmp_dir.name,
        )

        msg = f"Saved training job package to: {output_path}"
        print(msg)

        # Close training editor.
        self.accept()

        if gui:
            msgBox = QtWidgets.QMessageBox(text="Created training job package.")
            msgBox.setDetailedText(output_path)
            msgBox.setWindowTitle("Training Job Package")
            msgBox.addButton(QtWidgets.QMessageBox.Ok)
            openFolderButton = msgBox.addButton(
                "Open containing folder", QtWidgets.QMessageBox.ActionRole
            )
            colabButton = msgBox.addButton(
                "Go to Colab", QtWidgets.QMessageBox.ActionRole
            )
            msgBox.exec_()

            if msgBox.clickedButton() == openFolderButton:
                sleap.gui.commands.open_file(str(Path(output_path).resolve().parent))
            elif msgBox.clickedButton() == colabButton:
                # TODO: Update this to more workflow-tailored notebook.
                sleap.gui.commands.copy_to_clipboard(output_path)
                sleap.gui.commands.open_website(
                    "https://colab.research.google.com/github/talmolab/sleap/blob/main/docs/notebooks/Training_and_inference_using_Google_Drive.ipynb"
                )

        tmp_dir.cleanup()

frame_selection property writable

Returns dictionary with frames that user has selected for learning.

copy()

Copy scripts and configs to clipboard

Source code in sleap/gui/learning/dialog.py

def copy(self):
    """Copy scripts and configs to clipboard"""

    # Get all info from dialog
    pipeline_form_data = self.pipeline_form_widget.get_form_data()
    config_info_list = self.get_every_head_config_data(pipeline_form_data)

    # Format information for each tab in dialog
    # output = [OmegaConf.to_yaml(pipeline_form_data)] # TODO:cfg:
    output = []
    for config_info in config_info_list:
        config_info = config_info.config
        # convert to sleap-nn cfg (yaml)
        from sleap_nn.config.training_job_config import verify_training_cfg

        config_info = filter_cfg(config_info)
        cfg = verify_training_cfg(config_info)
        cfg.data_config.train_labels_path = [self.labels_filename]
        output.append(OmegaConf.to_yaml(cfg))

    output = "\n".join(output)
    # Set the clipboard text
    clipboard = QtWidgets.QApplication.clipboard()
    clipboard.setText(output)

export_package(output_path=None, gui=True)

Export training job package.

Source code in sleap/gui/learning/dialog.py

def export_package(self, output_path: Optional[str] = None, gui: bool = True):
    """Export training job package."""
    # TODO: Warn if self.mode != "training"?
    if output_path is None or not output_path:
        # Prompt for output path.
        output_path, _ = FileDialog.save(
            caption="Export Training Job Package...",
            dir=f"{self.labels_filename}.training_job.zip",
            filter="Training Job Package (*.zip)",
        )
        if len(output_path) == 0:
            return

    # Create temp dir before packaging.
    tmp_dir = tempfile.TemporaryDirectory()

    # Remove the temp dir when program exits in case something goes wrong.
    # atexit.register(shutil.rmtree, tmp_dir.name, ignore_errors=True)

    # Check if we need to include suggestions.
    include_suggestions = False
    items_for_inference = self.get_items_for_inference(
        self.pipeline_form_widget.get_form_data()
    )
    for item in items_for_inference.items:
        if (
            isinstance(item, runners.DatasetItemForInference)
            and item.frame_filter == "suggested"
        ):
            include_suggestions = True

    # Save dataset with images.
    labels_pkg_filename = str(
        Path(self.labels_filename).with_suffix(".pkg.slp").name
    )
    if gui:
        ret = sleap.gui.commands.export_dataset_gui(
            self.labels,
            tmp_dir.name + "/" + labels_pkg_filename,
            all_labeled=False,
            suggested=include_suggestions,
        )
        if ret == "canceled":
            # Quit if user canceled during export.
            tmp_dir.cleanup()
            return
    else:
        self.labels.save(
            filename=tmp_dir.name + "/" + labels_pkg_filename,
            embed=True,
        )

    # Save config and scripts.
    self.save(tmp_dir.name, labels_filename=labels_pkg_filename)

    # Package everything.
    shutil.make_archive(
        base_name=str(Path(output_path).with_suffix("")),
        format="zip",
        root_dir=tmp_dir.name,
    )

    msg = f"Saved training job package to: {output_path}"
    print(msg)

    # Close training editor.
    self.accept()

    if gui:
        msgBox = QtWidgets.QMessageBox(text="Created training job package.")
        msgBox.setDetailedText(output_path)
        msgBox.setWindowTitle("Training Job Package")
        msgBox.addButton(QtWidgets.QMessageBox.Ok)
        openFolderButton = msgBox.addButton(
            "Open containing folder", QtWidgets.QMessageBox.ActionRole
        )
        colabButton = msgBox.addButton(
            "Go to Colab", QtWidgets.QMessageBox.ActionRole
        )
        msgBox.exec_()

        if msgBox.clickedButton() == openFolderButton:
            sleap.gui.commands.open_file(str(Path(output_path).resolve().parent))
        elif msgBox.clickedButton() == colabButton:
            # TODO: Update this to more workflow-tailored notebook.
            sleap.gui.commands.copy_to_clipboard(output_path)
            sleap.gui.commands.open_website(
                "https://colab.research.google.com/github/talmolab/sleap/blob/main/docs/notebooks/Training_and_inference_using_Google_Drive.ipynb"
            )

    tmp_dir.cleanup()

run()

Run with current dialog settings.

Source code in sleap/gui/learning/dialog.py

def run(self):
    """Run with current dialog settings."""

    pipeline_form_data = self.pipeline_form_widget.get_form_data()

    items_for_inference = self.get_items_for_inference(pipeline_form_data)

    config_info_list = self.get_every_head_config_data(pipeline_form_data)

    # Close the dialog now that we have the data from it
    self.accept()

    # Run training/learning pipeline using the TrainingJobs
    new_counts = runners.run_learning_pipeline(
        labels_filename=self.labels_filename,
        labels=self.labels,
        config_info_list=config_info_list,
        inference_params=pipeline_form_data,
        items_for_inference=items_for_inference,
    )

    self._handle_learning_finished.emit(new_counts)

    # count < 0 means there was an error and we didn't get any results.
    if new_counts is not None and new_counts >= 0:
        total_count = items_for_inference.total_frame_count
        no_result_count = max(0, total_count - new_counts)

        message = (
            f"Inference ran on {total_count} frames."
            f"\n\nInstances were predicted on {new_counts} frames "
            f"({no_result_count} frame{'s' if no_result_count != 1 else ''} with "
            "no instances found)."
        )

        win = QtWidgets.QMessageBox(text=message)
        win.setWindowTitle("Inference Results")
        win.exec_()

save(output_dir=None, labels_filename=None)

Save scripts and configs to run pipeline.

Source code in sleap/gui/learning/dialog.py

def save(
    self, output_dir: Optional[str] = None, labels_filename: Optional[str] = None
):
    """Save scripts and configs to run pipeline."""
    if output_dir is None or not output_dir:
        labels_fn = Path(self.labels_filename)
        models_dir = Path(labels_fn.parent, "models")
        output_dir = FileDialog.openDir(
            None,
            dir=models_dir.as_posix(),
            caption="Select directory to save scripts",
        )

        if not output_dir:
            return

    pipeline_form_data = self.pipeline_form_widget.get_form_data()
    items_for_inference = self.get_items_for_inference(pipeline_form_data)
    config_info_list = self.get_every_head_config_data(pipeline_form_data)

    if labels_filename is None:
        labels_filename = self.labels_filename

    runners.write_pipeline_files(
        output_dir=output_dir,
        labels_filename=labels_filename,
        config_info_list=config_info_list,
        inference_params=pipeline_form_data,
        items_for_inference=items_for_inference,
    )

update_loaded_config(loaded_cfg, tab_cfg_key_val_dict) staticmethod

Update a loaded preset config with values from the training editor.

Parameters:

Name	Type	Description	Default
loaded_cfg	dict	Dict from a yaml file that was loaded from a preset or previous training run.	required
tab_cfg_key_val_dict	dict	A dictionary with the values extracted from the training editor GUI tab.	required

Returns:

Type	Description
	A ScopedKeyDict with the loaded config values overriden by the

corresponding ones from the tab_cfg_key_val_dict.

Source code in sleap/gui/learning/dialog.py

@staticmethod
def update_loaded_config(
    loaded_cfg: dict, tab_cfg_key_val_dict: dict
):  # -> scopedkeydict.ScopedKeyDict:
    """Update a loaded preset config with values from the training editor.

    Args:
        loaded_cfg: Dict from a yaml file that was loaded from a preset or previous
            training run.
        tab_cfg_key_val_dict: A dictionary with the values extracted from the
            training editor GUI tab.

    Returns:
                A `ScopedKeyDict` with the loaded config values overriden by the
    corresponding ones from the `tab_cfg_key_val_dict`.
    """
    # Replace params exposed in GUI with values from GUI
    for param, value in tab_cfg_key_val_dict.items():
        loaded_cfg[param] = value

    return loaded_cfg

TrainingEditorWidget

Bases: QWidget

Dialog for viewing and modifying training profiles (model hyperparameters).

Parameters:

Name	Type	Description	Default
video	Optional[Video]	Video to use for receptive field preview	None
skeleton	Optional[Skeleton]	Skeleton to use for node option list	None
head	Optional[Text]	If given, then only show configs with specified head name	None
cfg_getter	Optional[TrainingConfigsGetter]	Object to use for getting list of config files. If given, then menu of config files will be shown so user can either copy hyperameters from another profile/model, or use a model that was already trained.	None
require_trained	bool	If True, then only show configs that are trained, and don't allow user to uncheck "use trained" setting. This is set when :py:class:LearningDialog is in "inference" mode.	False

Source code in sleap/gui/learning/dialog.py

class TrainingEditorWidget(QtWidgets.QWidget):
    """
    Dialog for viewing and modifying training profiles (model hyperparameters).

    Args:
        video: `Video` to use for receptive field preview
        skeleton: `Skeleton` to use for node option list
        head: If given, then only show configs with specified head name
        cfg_getter: Object to use for getting list of config files.
            If given, then menu of config files will be shown so user can
            either copy hyperameters from another profile/model, or use a model
            that was already trained.
        require_trained: If True, then only show configs that are trained,
            and don't allow user to uncheck "use trained" setting. This is set
            when :py:class:`LearningDialog` is in "inference" mode.
    """

    valueChanged = QtCore.Signal()

    def __init__(
        self,
        video: Optional[Video] = None,
        skeleton: Optional["Skeleton"] = None,
        head: Optional[Text] = None,
        cfg_getter: Optional["TrainingConfigsGetter"] = None,
        require_trained: bool = False,
        *args,
        **kwargs,
    ):
        super(TrainingEditorWidget, self).__init__()

        self._video = video
        self._cfg_getter = cfg_getter
        self._cfg_list_widget = None
        self._receptive_field_widget = None
        self._use_trained_model = None
        self._resume_training = None
        self._require_trained = require_trained
        self.head = head

        yaml_name = "training_editor_form"

        self.form_widgets: Dict[str, YamlFormWidget] = dict()

        for key in ("model", "data", "augmentation", "optimization", "outputs"):
            self.form_widgets[key] = YamlFormWidget.from_name(
                yaml_name, which_form=key, title=key.title()
            )
            self.form_widgets[key].valueChanged.connect(self.emitValueChanged)

        self.form_widgets["model"].valueChanged.connect(self.update_receptive_field)
        self.form_widgets["data"].valueChanged.connect(self.update_receptive_field)

        if hasattr(skeleton, "node_names"):
            for field_name in NODE_LIST_FIELDS:
                form_name = field_name.split(".")[0]
                self.form_widgets[form_name].set_field_options(
                    ".".join(field_name.split(".")[1:]),
                    skeleton.node_names,
                )

        if self._video:
            self._receptive_field_widget = receptivefield.ReceptiveFieldWidget(
                self.head
            )
            self._receptive_field_widget.setImage(self._video.backend.read_test_frame())

        self._set_head()

        # Layout for header and columns
        layout = QtWidgets.QVBoxLayout()

        # Two column layout for config parameters
        col1_layout = QtWidgets.QVBoxLayout()
        col2_layout = QtWidgets.QVBoxLayout()
        col3_layout = QtWidgets.QVBoxLayout()

        col1_layout.addWidget(self.form_widgets["data"])
        col1_layout.addWidget(self.form_widgets["optimization"])
        col2_layout.addWidget(self.form_widgets["augmentation"])
        col3_layout.addWidget(self.form_widgets["model"])

        if self._receptive_field_widget:
            col0_layout = QtWidgets.QVBoxLayout()
            col0_layout.addWidget(self._receptive_field_widget)
        else:
            col0_layout = None

        col_layout = QtWidgets.QHBoxLayout()
        if col0_layout:
            col_layout.addWidget(self._layout_widget(col0_layout))
        col_layout.addWidget(self._layout_widget(col1_layout))
        col_layout.addWidget(self._layout_widget(col2_layout))
        col_layout.addWidget(self._layout_widget(col3_layout))

        # If we have an object which gets a list of config files,
        # then we'll show a menu to allow selection from the list.
        if self._cfg_getter is not None:
            self._cfg_list_widget = configs.TrainingConfigFilesWidget(
                cfg_getter=self._cfg_getter,
                head_name=cast(str, head),  # Expect head to be a string
                require_trained=require_trained,
            )
            self._cfg_list_widget.onConfigSelection.connect(
                self.acceptSelectedConfigInfo
            )
            # self._cfg_list_widget.setDataDict.connect(
            # self.set_fields_from_key_val_dict
            # )

            layout.addWidget(self._cfg_list_widget)

        if self._require_trained:
            self._update_use_trained()
        elif self._cfg_list_widget is not None:
            # Add option for using trained model from selected config file
            self._use_trained_model = QtWidgets.QCheckBox(
                "Use Existing Training Config"
            )
            self._use_trained_model.setEnabled(False)
            self._use_trained_model.setVisible(False)
            self._resume_training = QtWidgets.QCheckBox("Use Trained Model Weights")
            self._resume_training.setEnabled(False)
            self._resume_training.setVisible(False)

            self._use_trained_model.stateChanged.connect(self._update_use_trained)
            self._resume_training.stateChanged.connect(self._update_use_trained)

            layout.addWidget(self._use_trained_model)
            layout.addWidget(self._resume_training)

        layout.addWidget(self._layout_widget(col_layout))
        self.setLayout(layout)

    @classmethod
    def from_trained_config(
        cls, cfg_info: configs.ConfigFileInfo, cfg_getter: configs.TrainingConfigsGetter
    ):
        widget = cls(
            require_trained=True, head=cfg_info.head_name, cfg_getter=cfg_getter
        )
        widget.acceptSelectedConfigInfo(cfg_info)
        widget.setWindowTitle(cfg_info.path_dir)
        return widget

    @staticmethod
    def _layout_widget(layout):
        widget = QtWidgets.QWidget()
        widget.setLayout(layout)
        return widget

    def emitValueChanged(self):
        self.valueChanged.emit()

        # When there's a config getter, we want to inform it that the data
        # has changed so that it can activate/update the "user" config
        # if self._cfg_list_widget:
        #     self._set_user_config()

    def acceptSelectedConfigInfo(self, cfg_info: configs.ConfigFileInfo):
        self._load_config(cfg_info)

        has_trained_model = cfg_info.has_trained_model
        if self._use_trained_model is not None:
            self._use_trained_model.setChecked(self._require_trained)
            self._use_trained_model.setVisible(has_trained_model)
            self._use_trained_model.setEnabled(has_trained_model)
        # Redundant check (for readability) since this checkbox exists if the
        # above does
        if self._resume_training is not None:
            self._use_trained_model.setChecked(False)
            self._resume_training.setVisible(has_trained_model)
            self._resume_training.setEnabled(has_trained_model)

        self.update_receptive_field()

    def update_receptive_field(self):
        data_form_data = get_omegaconf_from_gui_form(
            self.form_widgets["data"].get_form_data()
        )

        model_cfg = get_omegaconf_from_gui_form(
            self.form_widgets["model"].get_form_data()
        )

        rf_image_scale = OmegaConf.select(
            data_form_data, "data_config.preprocessing.scale", default=1.0
        )

        if self._receptive_field_widget:
            self._receptive_field_widget.setModelConfig(model_cfg, scale=rf_image_scale)
            self._receptive_field_widget.repaint()

    def update_file_list(self):
        self._cfg_list_widget.update()

    def _load_config_or_key_val_dict(self, cfg_data):
        if type(cfg_data) != dict:
            self._load_config(cfg_data)
        else:
            self.set_fields_from_key_val_dict(cfg_data)

    def _load_config(self, cfg_info: configs.ConfigFileInfo):
        if cfg_info is None:
            return

        cfg = cfg_info.config
        key_val_dict = get_keyval_dict_from_omegaconf(cfg)
        if key_val_dict.get("trainer_config.trainer_devices") == "auto":
            key_val_dict["trainer_config.trainer_devices"] = None
        self.set_fields_from_key_val_dict(key_val_dict)

    # def _set_user_config(self):
    #     cfg_form_data_dict = self.get_all_form_data()
    #     self._cfg_list_widget.setUserConfigData(cfg_form_data_dict)

    def _update_use_trained(self, check_state=0):
        """Update config GUI based on _use_trained_model & _resume_training checkboxes.

        This function is called when either _use_trained_model or _resume_training
        checkbox is checked/unchecked or when _require_trained is changed.

        If _require_trained is True, then we'll disable all fields.
        If _use_trained_model is checked, then we'll disable all fields.
        If _resume_training is checked, then we'll disable only the model field.

        Args:
            check_state (int, optional): Check state of checkbox. Defaults to 0. Unused.

        Returns:
            None

        Side Effects:
            Disables/Enables fields based on checkbox values
            (and _required_training).
        """

        # Check which checkbox changed its value (if any)
        sender = self.sender()

        if sender is None:  # If sender is None, then _required_training is True
            pass
        # Uncheck _resume_training checkbox if _use_trained_model is unchecked
        elif (
            sender == self._use_trained_model
            and not self._use_trained_model.isChecked()
        ):
            self._resume_training.setChecked(False)

        # Check _use_trained_model checkbox if _resume_training is checked
        elif sender == self._resume_training and self._resume_training.isChecked():
            self._use_trained_model.setChecked(True)

        # Update form widgets
        use_trained_params = self.use_trained
        use_model_params = self.resume_training
        for form in self.form_widgets.values():
            form.set_enabled(not use_trained_params)

        if use_trained_params or use_model_params:
            cfg_info = self._cfg_list_widget.getSelectedConfigInfo()

        # If user wants to resume training, then reset only model form to match config
        if use_model_params:
            self.form_widgets["model"].set_enabled(False)

            # Set model form to match config
            cfg = cfg_info.config
            key_val_dict = get_keyval_dict_from_omegaconf(cfg)
            self.set_fields_from_key_val_dict({"model": key_val_dict})

        # If user wants to use trained model, then reset entire form to match config
        if use_trained_params:
            self._load_config(cfg_info)

        self._set_head()

    def _set_head(self):
        if self.head:
            self.set_fields_from_key_val_dict(
                {
                    "_heads_name": self.head,
                }
            )

            self.form_widgets["model"].set_field_enabled("_heads_name", False)

    def set_fields_from_key_val_dict(self, cfg_key_val_dict):
        for form in self.form_widgets.values():
            form.set_form_data(cfg_key_val_dict)

        self._set_backbone_from_key_val_dict(cfg_key_val_dict)

    def _set_backbone_from_key_val_dict(self, cfg_key_val_dict):
        for key, val in cfg_key_val_dict.items():
            if (
                key.startswith("model.model_config.backbone_config.")
                and val is not None
            ):
                backbone_name = key.split(".")[3]
                self.set_fields_from_key_val_dict(dict(_backbone_name=backbone_name))
                break

    @property
    def use_trained(self) -> bool:
        if self._require_trained or (
            (self._use_trained_model is not None)
            and self._use_trained_model.isChecked()
            and (not self.resume_training)
        ):
            return True

        return False

    @property
    def resume_training(self) -> bool:
        if (self._resume_training is not None) and self._resume_training.isChecked():
            return True
        return False

    @property
    def trained_config_info_to_use(self) -> Optional[configs.ConfigFileInfo]:
        # If `TrainingEditorWidget` was initialized with a config getter, then
        # we expect to have a list of config files
        if self._cfg_list_widget is None:
            return None

        selected_config_info: Optional[configs.ConfigFileInfo] = (
            self._cfg_list_widget.getSelectedConfigInfo()
        )
        if (selected_config_info is None) or (
            not selected_config_info.has_trained_model
        ):
            return None

        trained_config_info = configs.ConfigFileInfo.from_config_file(
            selected_config_info.path
        )
        if self.use_trained:
            trained_config_info.dont_retrain = True
        else:
            # Set certain parameters to defaults
            trained_config = trained_config_info.config
            trained_config.data_config.val_labels_path = None
            trained_config.data_config.test_file_path = None
            trained_config.data_config.skeletons = []
            trained_config.trainer_config.ckpt_dir = None
            trained_config.trainer_config.run_name = None

        if self.resume_training:
            # Get the folder path of trained config and set it as the output
            # folder
            file_list = list(Path(cast(str, trained_config_info.path)).parent.iterdir())
            if (
                Path(cast(str, trained_config_info.path)).parent / "best.ckpt"
            ) in file_list:
                ckpt = "best.ckpt"
            elif (
                Path(cast(str, trained_config_info.path)).parent / "best_model.h5"
            ) in file_list:
                ckpt = "best_model.h5"
            trained_config_info.config.model_config.pretrained_backbone_weights = (
                Path(cast(str, trained_config_info.path)).parent / ckpt
            ).as_posix()
            trained_config_info.config.model_config.pretrained_head_weights = (
                trained_config_info.config.model_config.pretrained_backbone_weights
            )
        else:
            trained_config_info.config.model_config.pretrained_backbone_weights = None
            trained_config_info.config.model_config.pretrained_head_weights = None

        return trained_config_info

    @property
    def has_trained_config_selected(self) -> bool:
        if self._cfg_list_widget is None:
            return False

        cfg_info = self._cfg_list_widget.getSelectedConfigInfo()
        if cfg_info and cfg_info.has_trained_model:
            return True

        return False

    def get_all_form_data(self) -> dict:
        form_data = dict()
        for form in self.form_widgets.values():
            form_data.update(form.get_form_data())
        return form_data

TrainingPipelineWidget

Bases: QWidget

Widget used in :py:class:LearningDialog for configuring pipeline.

Source code in sleap/gui/learning/dialog.py

class TrainingPipelineWidget(QtWidgets.QWidget):
    """
    Widget used in :py:class:`LearningDialog` for configuring pipeline.
    """

    updatePipeline = QtCore.Signal(str)
    valueChanged = QtCore.Signal()

    def __init__(
        self, mode: Text, skeleton: Optional["Skeleton"] = None, *args, **kwargs
    ):
        super(TrainingPipelineWidget, self).__init__(*args, **kwargs)

        self.form_widget = YamlFormWidget.from_name(
            "pipeline_form", which_form=mode, title="Training Pipeline"
        )

        if hasattr(skeleton, "node_names"):
            for field_name in NODE_LIST_FIELDS:
                self.form_widget.set_field_options(
                    ".".join(field_name.split(".")[1:]),
                    skeleton.node_names,
                )

        # Connect actions for change to pipeline
        self.pipeline_field = self.form_widget.form_layout.find_field("_pipeline")[0]
        self.pipeline_field.valueChanged.connect(self.emitPipeline)

        self.form_widget.form_layout.valueChanged.connect(self.valueChanged)

        self.setLayout(self.form_widget.form_layout)

    @property
    def fields(self):
        return self.form_widget.fields

    @property
    def buttons(self):
        return self.form_widget.buttons

    def set_message(self, message: Text):
        self.form_widget.set_message()

    def get_form_data(self):
        return self.form_widget.get_form_data()

    def set_form_data(self, data):
        self.form_widget.set_form_data(data)

    def emitPipeline(self):
        val = self.current_pipeline
        self.updatePipeline.emit(val)

    @property
    def current_pipeline(self):
        pipeline_selected_label = self.pipeline_field.value()
        if "top-down" in pipeline_selected_label:
            if "id" not in pipeline_selected_label:
                return "top-down"
            else:
                return "top-down-id"
        if "bottom-up" in pipeline_selected_label:
            if "id" not in pipeline_selected_label:
                return "bottom-up"
            else:
                return "bottom-up-id"
        if "single" in pipeline_selected_label:
            return "single"
        return ""

    @current_pipeline.setter
    def current_pipeline(self, val):
        if val not in (
            "top-down",
            "bottom-up",
            "single",
            "top-down-id",
            "bottom-up-id",
        ):
            raise ValueError(f"Cannot set pipeline to {val}")

        # Match short name to full pipeline name shown in menu
        for full_option_name in self.pipeline_field.option_list:
            if val in full_option_name:
                val = full_option_name
                break

        self.pipeline_field.setValue(val)
        self.emitPipeline()