def from_trained_models(
cls,
centroid_model_path: Optional[Text] = None,
confmap_model_path: Optional[Text] = None,
batch_size: int = 1,
peak_threshold: float = 0.2,
integral_refinement: bool = True,
integral_patch_size: int = 5,
) -> "TopdownPredictor":
"""Create predictor from saved models.
Args:
centroid_model_path: Path to centroid model folder.
confmap_model_path: Path to topdown confidence map model folder.
Returns:
An instance of TopdownPredictor with the loaded models.
One of the two models can be left as None to perform inference with ground
truth data. This will only work with LabelsReader as the provider.
"""
if centroid_model_path is None and confmap_model_path is None:
raise ValueError(
"Either the centroid or topdown confidence map model must be provided."
)
if centroid_model_path is not None:
# Load centroid model.
centroid_config = TrainingJobConfig.load_json(centroid_model_path)
centroid_keras_model_path = get_keras_model_path(centroid_model_path)
centroid_model = Model.from_config(centroid_config.model)
centroid_model.keras_model = tf.keras.models.load_model(
centroid_keras_model_path, compile=False
)
else:
centroid_config = None
centroid_model = None
if confmap_model_path is not None:
# Load confmap model.
confmap_config = TrainingJobConfig.load_json(confmap_model_path)
confmap_keras_model_path = get_keras_model_path(confmap_model_path)
confmap_model = Model.from_config(confmap_config.model)
confmap_model.keras_model = tf.keras.models.load_model(
confmap_keras_model_path, compile=False
)
else:
confmap_config = None
confmap_model = None
return cls(
centroid_config=centroid_config,
centroid_model=centroid_model,
confmap_config=confmap_config,
confmap_model=confmap_model,
batch_size=batch_size,
peak_threshold=peak_threshold,
integral_refinement=integral_refinement,
integral_patch_size=integral_patch_size,
)
def from_trained_models(cls, model_path: Text) -> "VisualPredictor":
cfg = TrainingJobConfig.load_json(model_path)
keras_model_path = get_keras_model_path(model_path)
model = Model.from_config(cfg.model)
model.keras_model = tf.keras.models.load_model(keras_model_path, compile=False)
return cls(config=cfg, model=model)
def from_config_file(cls, path):
cfg = TrainingJobConfig.load_json(path)
head_name = cfg.model.heads.which_oneof_attrib_name()
filename = os.path.basename(path)
return cls(config=cfg,
path=path,
filename=filename,
head_name=head_name)
def from_trained_models(cls, bottomup_model_path: Text) -> "BottomupPredictor":
"""Create predictor from saved models."""
# Load bottomup model.
bottomup_config = TrainingJobConfig.load_json(bottomup_model_path)
bottomup_keras_model_path = get_keras_model_path(bottomup_model_path)
bottomup_model = Model.from_config(bottomup_config.model)
bottomup_model.keras_model = tf.keras.models.load_model(
bottomup_keras_model_path, compile=False
)
return cls(bottomup_config=bottomup_config, bottomup_model=bottomup_model)
def try_loading_path(self, path: Text):
try:
cfg = TrainingJobConfig.load_json(path)
except Exception as e:
# Couldn't load so just ignore
print(e)
pass
else:
# Get the head from the model (i.e., what the model will predict)
key = cfg.model.heads.which_oneof_attrib_name()
filename = os.path.basename(path)
# If filter isn't set or matches head name, add config to list
if self.head_filter in (None, key):
return ConfigFileInfo(path=path,
filename=filename,
config=cfg,
head_name=key)
return None
def find_heads_for_model_paths(paths) -> Dict[str, str]:
"""Given list of models paths, returns dict with path keyed by head name."""
trained_model_paths = dict()
if paths is None:
return trained_model_paths
for model_path in paths:
# Load the model config
cfg = TrainingJobConfig.load_json(model_path)
# Get the head from the model (i.e., what the model will predict)
key = cfg.model.heads.which_oneof_attrib_name()
# If path is to config file json, then get the path to parent dir
if model_path.endswith(".json"):
model_path = os.path.dirname(model_path)
trained_model_paths[key] = model_path
return trained_model_paths
def from_trained_models(
cls,
confmap_model_path: Text,
peak_threshold: float = 0.2,
integral_refinement: bool = True,
integral_patch_size: int = 7,
) -> "SingleInstancePredictor":
"""Create predictor from saved models."""
# Load confmap model.
confmap_config = TrainingJobConfig.load_json(confmap_model_path)
confmap_keras_model_path = get_keras_model_path(confmap_model_path)
confmap_model = Model.from_config(confmap_config.model)
confmap_model.keras_model = tf.keras.models.load_model(
confmap_keras_model_path, compile=False)
return cls(
confmap_config=confmap_config,
confmap_model=confmap_model,
peak_threshold=peak_threshold,
integral_refinement=integral_refinement,
integral_patch_size=integral_patch_size,
)
def train_subprocess(
job_config: TrainingJobConfig,
labels_filename: str,
video_paths: Optional[List[Text]] = None,
waiting_callback: Optional[Callable] = None,
save_viz: bool = False,
):
"""Runs training inside subprocess."""
# run_name = job_config.outputs.run_name
run_path = job_config.outputs.run_path
success = False
with tempfile.TemporaryDirectory() as temp_dir:
# Write a temporary file of the TrainingJob so that we can respect
# any changed made to the job attributes after it was loaded.
temp_filename = datetime.now().strftime(
"%y%m%d_%H%M%S") + "_training_job.json"
training_job_path = os.path.join(temp_dir, temp_filename)
job_config.save_json(training_job_path)
# Build CLI arguments for training
cli_args = [
"sleap-train",
training_job_path,
labels_filename,
"--zmq",
]
if save_viz:
cli_args.append("--save_viz")
# Use cli arg since cli ignores setting in config
if job_config.outputs.tensorboard.write_logs:
cli_args.append("--tensorboard")
# Add list of video paths so we can find video even if paths in saved
# labels dataset file are incorrect.
if video_paths:
cli_args.extend(("--video-paths", ",".join(video_paths)))
print(cli_args)
if not SKIP_TRAINING:
# Run training in a subprocess
with sub.Popen(cli_args) as proc:
# Wait till training is done, calling a callback if given.
while proc.poll() is None:
if waiting_callback is not None:
if waiting_callback() == -1:
# -1 signals user cancellation
return "", False
time.sleep(0.1)
success = proc.returncode == 0
print("Run Path:", run_path)
return run_path, success
def main():
"""Create CLI for training and run."""
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("training_job_path",
help="Path to training job profile JSON file.")
parser.add_argument("labels_path",
help="Path to labels file to use for training.")
parser.add_argument(
"--video-paths",
type=str,
default="",
help=
"List of paths for finding videos in case paths inside labels file need fixing.",
)
parser.add_argument(
"--val_labels",
"--val",
help=
"Path to labels file to use for validation (overrides training job path if set).",
)
parser.add_argument(
"--test_labels",
"--test",
help=
"Path to labels file to use for test (overrides training job path if set).",
)
parser.add_argument(
"--tensorboard",
action="store_true",
help="Enables TensorBoard logging to the run path.",
)
parser.add_argument(
"--save_viz",
action="store_true",
help="Enables saving of prediction visualizations to the run folder.",
)
parser.add_argument("--zmq",
action="store_true",
help="Enables ZMQ logging (for GUI).")
parser.add_argument(
"--run_name",
default="",
help=
"Run name to use when saving file, overrides other run name settings.",
)
parser.add_argument("--prefix",
default="",
help="Prefix to prepend to run name.")
parser.add_argument("--suffix",
default="",
help="Suffix to append to run name.")
args, _ = parser.parse_known_args()
# Find job configuration file.
job_filename = args.training_job_path
if not os.path.exists(job_filename):
profile_dir = get_package_file("sleap/training_profiles")
if os.path.exists(os.path.join(profile_dir, job_filename)):
job_filename = os.path.join(profile_dir, job_filename)
else:
raise FileNotFoundError(
f"Could not find training profile: {job_filename}")
# Load job configuration.
job_config = TrainingJobConfig.load_json(job_filename)
# Override config settings for CLI-based training.
job_config.outputs.save_outputs = True
job_config.outputs.tensorboard.write_logs = args.tensorboard
job_config.outputs.zmq.publish_updates = args.zmq
job_config.outputs.zmq.subscribe_to_controller = args.zmq
if args.run_name != "":
job_config.outputs.run_name = args.run_name
if args.prefix != "":
job_config.outputs.run_name_prefix = args.prefix
if args.suffix != "":
job_config.outputs.run_name_suffix = args.suffix
job_config.outputs.save_visualizations = args.save_viz
logger.info(f"Training labels file: {args.labels_path}")
logger.info(f"Training profile: {job_filename}")
logger.info("")
# Log configuration to console.
logger.info("Arguments:")
logger.info(json.dumps(vars(args), indent=4))
logger.info("")
logger.info("Training job:")
logger.info(job_config.to_json())
logger.info("")
logger.info("Initializing trainer...")
# Create a trainer and run!
trainer = Trainer.from_config(
job_config,
training_labels=args.labels_path,
validation_labels=args.val_labels,
test_labels=args.test_labels,
video_search_paths=args.video_paths.split(","),
)
trainer.train()