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_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 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 receptive_field_info_from_model_cfg(model_cfg: ModelConfig) -> dict:
"""Gets receptive field information given specific model configuration."""
rf_info = dict(
size=None,
max_stride=None,
down_blocks=None,
convs_per_block=None,
kernel_size=None,
)
try:
model = Model.from_config(model_cfg, Skeleton())
except ZeroDivisionError:
# Unable to create model from these config parameters
return rf_info
if hasattr(model_cfg.backbone.which_oneof(), "max_stride"):
rf_info["max_stride"] = model_cfg.backbone.which_oneof().max_stride
if hasattr(model.backbone, "down_convs_per_block"):
rf_info["convs_per_block"] = model.backbone.down_convs_per_block
elif hasattr(model.backbone, "convs_per_block"):
rf_info["convs_per_block"] = model.backbone.convs_per_block
if hasattr(model.backbone, "kernel_size"):
rf_info["kernel_size"] = model.backbone.kernel_size
rf_info["down_blocks"] = model.backbone.down_blocks
if rf_info["down_blocks"] and rf_info["convs_per_block"] and rf_info[
"kernel_size"]:
rf_info["size"] = compute_rf(
down_blocks=rf_info["down_blocks"],
convs_per_block=rf_info["convs_per_block"],
kernel_size=rf_info["kernel_size"],
)
return rf_info
def from_config(
cls,
config: TrainingJobConfig,
training_labels: Optional[Union[Text, sleap.Labels]] = None,
validation_labels: Optional[Union[Text, sleap.Labels, float]] = None,
test_labels: Optional[Union[Text, sleap.Labels]] = None,
video_search_paths: Optional[List[Text]] = None,
) -> "Trainer":
"""Initialize the trainer from a training job configuration.
Args:
config: A `TrainingJobConfig` instance.
training_labels: Training labels to use instead of the ones in the config,
if any. If a path is specified, it will overwrite the one in the config.
validation_labels: Validation labels to use instead of the ones in the
config, if any. If a path is specified, it will overwrite the one in
the config.
test_labels: Teset labels to use instead of the ones in the config, if any.
If a path is specified, it will overwrite the one in the config.
"""
# Copy input config before we make any changes.
initial_config = copy.deepcopy(config)
# Create data readers and store loaded skeleton.
data_readers = DataReaders.from_config(
config.data.labels,
training=training_labels,
validation=validation_labels,
test=test_labels,
video_search_paths=video_search_paths,
update_config=True,
)
config.data.labels.skeletons = data_readers.training_labels.skeletons
# Create model.
model = Model.from_config(config.model,
skeleton=config.data.labels.skeletons[0],
update_config=True)
# Determine output type to create type-specific model trainer.
head_config = config.model.heads.which_oneof()
trainer_cls = None
if isinstance(head_config, CentroidsHeadConfig):
trainer_cls = CentroidConfmapsModelTrainer
elif isinstance(head_config, CenteredInstanceConfmapsHeadConfig):
trainer_cls = TopdownConfmapsModelTrainer
elif isinstance(head_config, MultiInstanceConfig):
trainer_cls = BottomUpModelTrainer
elif isinstance(head_config, SingleInstanceConfmapsHeadConfig):
trainer_cls = SingleInstanceModelTrainer
else:
raise ValueError(
"Model head not specified or configured. Check the config.model.heads"
" setting.")
return trainer_cls(
config=config,
initial_config=initial_config,
data_readers=data_readers,
model=model,
)