def run_gui_training(
labels_filename: str,
labels: Labels,
config_info_list: List[ConfigFileInfo],
gui: bool = True,
save_viz: bool = False,
) -> Dict[Text, Text]:
"""
Runs training for each training job.
Args:
labels: Labels object from which we'll get training data.
config_info_list: List of ConfigFileInfo with configs for training.
gui: Whether to show gui windows and process gui events.
save_viz: Whether to save visualizations from training.
Returns:
Dictionary, keys are head name, values are path to trained config.
"""
trained_job_paths = dict()
if gui:
from sleap.nn.monitor import LossViewer
from sleap.gui.widgets.imagedir import QtImageDirectoryWidget
# open training monitor window
win = LossViewer()
win.resize(600, 400)
win.show()
for config_info in config_info_list:
if config_info.dont_retrain:
if not config_info.has_trained_model:
raise ValueError(
f"Config is set to not retrain but no trained model found: {config_info.path}"
)
print(
f"Using already trained model for {config_info.head_name}: {config_info.path}"
)
trained_job_paths[config_info.head_name] = config_info.path
else:
job = config_info.config
model_type = config_info.head_name
# We'll pass along the list of paths we actually used for loading
# the videos so that we don't have to rely on the paths currently
# saved in the labels file for finding videos.
video_path_list = [video.filename for video in labels.videos]
# Update save dir and run name for job we're about to train
# so we have access to them here (rather than letting
# train_subprocess update them).
# training.Trainer.set_run_name(job, labels_filename)
job.outputs.runs_folder = os.path.join(
os.path.dirname(labels_filename), "models")
training.setup_new_run_folder(
job.outputs, base_run_name=f"{model_type}.{len(labels)}")
if gui:
print("Resetting monitor window.")
win.reset(what=str(model_type))
win.setWindowTitle(f"Training Model - {str(model_type)}")
win.set_message(f"Preparing to run training...")
if save_viz:
viz_window = QtImageDirectoryWidget.make_training_vizualizer(
job.outputs.run_path)
viz_window.move(win.x() + win.width() + 20, win.y())
win.on_epoch.connect(viz_window.poll)
print(f"Start training {str(model_type)}...")
def waiting():
if gui:
QtWidgets.QApplication.instance().processEvents()
# Run training
trained_job_path, success = train_subprocess(
job_config=job,
labels_filename=labels_filename,
video_paths=video_path_list,
waiting_callback=waiting,
save_viz=save_viz,
)
if success:
# get the path to the resulting TrainingJob file
trained_job_paths[model_type] = trained_job_path
print(f"Finished training {str(model_type)}.")
else:
if gui:
win.close()
QtWidgets.QMessageBox(
text=
f"An error occurred while training {str(model_type)}. Your command line terminal may have more information about the error."
).exec_()
trained_job_paths[model_type] = None
if gui:
# close training monitor window
win.close()
return trained_job_paths
def write_pipeline_files(
output_dir: str,
labels_filename: str,
config_info_list: List[ConfigFileInfo],
inference_params: Dict[str, Any],
items_for_inference: ItemsForInference,
):
"""Writes the config files and scripts for manually running pipeline."""
# Use absolute path for all files that aren't contained in the output dir.
labels_filename = os.path.abspath(labels_filename)
# Preserve current working directory and change working directory to the
# output directory, so we can set local paths relative to that.
old_cwd = os.getcwd()
os.chdir(output_dir)
new_cfg_filenames = []
train_script = "#!/bin/bash\n"
for cfg_info in config_info_list:
if cfg_info.dont_retrain:
# Use full absolute path to already training model
trained_path = os.path.normpath(
os.path.join(old_cwd, cfg_info.path))
new_cfg_filenames.append(trained_path)
else:
# We're training this model, so save config file...
# First we want to set the run folder so that we know where to find
# the model after it's trained.
# We'll use local path to the output directory (cwd).
# Note that setup_new_run_folder does things relative to cwd which
# is the main reason we're setting it to the output directory rather
# than just using normpath.
cfg_info.config.outputs.runs_folder = ""
training.setup_new_run_folder(cfg_info.config.outputs)
# Now we set the filename for the training config file
new_cfg_filename = f"{cfg_info.head_name}.json"
# Save the config file
cfg_info.config.save_json(new_cfg_filename)
# Keep track of the path where we'll find the trained model
new_cfg_filenames.append(cfg_info.config.outputs.run_name)
# Add a line to the script for training this model
train_script += f"sleap-train {new_cfg_filename} {labels_filename}\n"
# Write the script to train the models which need to be trained
with open(os.path.join(output_dir, "train-script.sh"), "w") as f:
f.write(train_script)
# Build the script for running inference
inference_script = "#!/bin/bash\n"
# Object with settings for inference
inference_task = InferenceTask(
labels_filename=labels_filename,
trained_job_paths=new_cfg_filenames,
inference_params=inference_params,
)
for item_for_inference in items_for_inference.items:
# We want to save predictions in output dir so use local path
prediction_output_path = (
f"{os.path.basename(item_for_inference.path)}.predictions.slp")
# Use absolute path to video
item_for_inference.use_absolute_path = True
# Get list of cli args
cli_args, _ = inference_task.make_predict_cli_call(
item_for_inference=item_for_inference,
output_path=prediction_output_path,
)
# And join them into a single call to inference
inference_script += " ".join(cli_args) + "\n"
# And write it
with open(os.path.join(output_dir, "inference-script.sh"), "w") as f:
f.write(inference_script)
# Restore the working directory
os.chdir(old_cwd)