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(
" –> ".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)
try:
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))
except ImportError:
show_sleap_nn_installation_message()
print(
"sleap-nn is not installed. This appears to be GUI-only install."
"To enable training, please install SLEAP with the 'nn' dependency."
"See the installation guide: https://docs.sleap.ai/latest/installation/"
)
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()
class: