def load_model_assets(
model_paths: List[str],
*,
device: str = "cpu",
backbone_ckpt_path: Optional[str] = None,
head_ckpt_path: Optional[str] = None,
peak_threshold: Union[float, List[float]] = 0.2,
integral_refinement: str = "integral",
integral_patch_size: int = 5,
max_instances: Optional[int] = None,
return_confmaps: bool = False,
preprocess_config: Optional["DictConfig"] = None,
anchor_part: Optional[str] = None,
max_edge_length_ratio: float = 0.25,
dist_penalty_weight: float = 1.0,
n_points: int = 10,
min_instance_peaks: Union[int, float] = 0,
min_line_scores: float = 0.25,
fg_threshold: float = 0.5,
min_mask_area: int = 0,
center_nms_kernel: int = 3,
mask_cleanup: bool = False,
mask_cleanup_radius: int = 0,
distance_gate_alpha: Optional[float] = None,
merge_fragments: bool = False,
merge_method: str = "greedy",
merge_thresholds: tuple = (0.85, 0.6, 0.4),
merge_w_valley: float = 1.0,
merge_w_offset: float = 0.25,
merge_dilate: int = 1,
full_res_masks: bool = False,
mask_output: str = "mask",
polygon_epsilon: float = 0.01,
bbox_nms_iou: Optional[float] = None,
) -> tuple[LoadedAssets, List[str]]:
"""Load checkpoints and build inference models.
Each entry in ``model_paths`` may be a model directory, or a path to its
``best.ckpt`` checkpoint or ``training_config.{yaml,json}`` config file; every
form is resolved to the model directory (#575).
The five bottom-up PAF grouping knobs (``max_edge_length_ratio``,
``dist_penalty_weight``, ``n_points``, ``min_instance_peaks``,
``min_line_scores``) are forwarded ONLY to the plain bottom-up builder
(legacy applied them only to ``BottomUpPredictor``); they are inert for
other model types. Likewise ``fg_threshold`` / ``min_mask_area`` are
forwarded ONLY to the bottom-up segmentation builder, and ``bbox_nms_iou``
only to the ``bbox_*`` detector builder.
Returns:
``(loaded_assets, model_types)`` — *model_types* is the list of
detected model types (one per path in *model_paths*).
"""
if preprocess_config is None:
preprocess_config = OmegaConf.create(
{
"ensure_rgb": None,
"ensure_grayscale": None,
"crop_size": None,
"max_width": None,
"max_height": None,
"scale": None,
}
)
# Accept a model directory, a best.ckpt path, or a training_config.{yaml,json}
# path for each entry; resolve every form to its model directory before
# detection/dispatch (the builders below index back into model_paths and join
# `best.ckpt` onto it). #575.
model_paths = [resolve_model_dir(mp) for mp in model_paths]
model_types: List[str] = []
configs: List[Any] = []
for mp in model_paths:
cfg, _ = _load_training_config(mp)
configs.append(cfg)
model_types.append(get_model_type_from_cfg(config=cfg))
# Reject duplicate model types up front. The dispatch below picks the
# FIRST path of a given type via `model_types.index(...)`, so passing two
# paths of the same type (e.g. two centroid dirs) would otherwise
# silently use only one and drop the other with no indication anything
# was wrong. (Detecting an *unrelated* extra path -- one whose type isn't
# consumed by whichever branch below ends up winning -- would need
# dispatch-branch-aware validation; that's a separate, bigger follow-up,
# not covered here.)
_seen_type_paths: dict = {}
for mp, mt in zip(model_paths, model_types):
if mt in _seen_type_paths:
raise ValueError(
f"Duplicate model type {mt!r} in --model_paths: got both "
f"{_seen_type_paths[mt]!r} and {mp!r}. Pass only one model "
"directory per type."
)
_seen_type_paths[mt] = mp
common_kwargs = dict(
device=device,
backbone_ckpt_path=backbone_ckpt_path,
head_ckpt_path=head_ckpt_path,
peak_threshold=peak_threshold,
integral_refinement=integral_refinement,
integral_patch_size=integral_patch_size,
return_confmaps=return_confmaps,
preprocess_config=preprocess_config,
)
# Dispatch on detected model type. The order (bottomup checked before the
# topdown family) differs from legacy but is inert for parity: each model
# path is exactly one type and bottom-up is a single-stage model never
# combined with a centroid/centered-instance pair, so the branches are
# mutually exclusive (a mixed bottomup+topdown path list is not a supported
# workflow and falls through to the final ValueError). #584.
if "single_instance" in model_types:
path = model_paths[model_types.index("single_instance")]
assets = _build_single_instance(path, **common_kwargs)
elif "bottomup" in model_types:
path = model_paths[model_types.index("bottomup")]
assets = _build_bottomup(
path,
max_instances=max_instances,
max_edge_length_ratio=max_edge_length_ratio,
dist_penalty_weight=dist_penalty_weight,
n_points=n_points,
min_instance_peaks=min_instance_peaks,
min_line_scores=min_line_scores,
**common_kwargs,
)
elif "multi_class_bottomup" in model_types:
path = model_paths[model_types.index("multi_class_bottomup")]
assets = _build_bottomup_multiclass(
path, max_instances=max_instances, **common_kwargs
)
elif "bottomup_segmentation" in model_types:
path = model_paths[model_types.index("bottomup_segmentation")]
assets = _build_bottomup_segmentation(
path,
fg_threshold=fg_threshold,
min_mask_area=min_mask_area,
max_instances=max_instances,
center_nms_kernel=center_nms_kernel,
mask_cleanup=mask_cleanup,
mask_cleanup_radius=mask_cleanup_radius,
distance_gate_alpha=distance_gate_alpha,
merge_fragments=merge_fragments,
merge_method=merge_method,
merge_thresholds=merge_thresholds,
merge_w_valley=merge_w_valley,
merge_w_offset=merge_w_offset,
merge_dilate=merge_dilate,
full_res_masks=full_res_masks,
mask_output=mask_output,
polygon_epsilon=polygon_epsilon,
**common_kwargs,
)
elif "semantic_segmentation" in model_types:
# Whole-frame semantic (foreground/background) segmentation. A single-stage
# whole-frame model (never combined with a centroid), so it gets its own
# top-level branch alongside bottomup_segmentation.
path = model_paths[model_types.index("semantic_segmentation")]
assets = _build_semantic_segmentation(
path,
fg_threshold=fg_threshold,
min_mask_area=min_mask_area,
full_res_masks=full_res_masks,
mask_output=mask_output,
polygon_epsilon=polygon_epsilon,
**common_kwargs,
)
elif "bbox_centernet" in model_types or "bbox_ssd" in model_types:
# Whole-frame bounding-box detectors. Alone they are single-stage; with a
# crop-consuming top-down model they replace the centroid model as stage 1.
model_type = "bbox_centernet" if "bbox_centernet" in model_types else "bbox_ssd"
if "bbox_centernet" in model_types and "bbox_ssd" in model_types:
raise ValueError(
"Got two bounding-box detectors (bbox_centernet and bbox_ssd); pass "
"one."
)
path = model_paths[model_types.index(model_type)]
stage_2 = [t for t in BBOX_STAGE_2_MODEL_TYPES if t in model_types]
detector_kwargs = dict(common_kwargs)
# `_resolve_preprocess_config` fills its argument in place: give the
# detector its own copy of the caller's overrides, so its training values
# (input scale, size matcher) do not leak into stage 2's resolution.
detector_kwargs["preprocess_config"] = copy.deepcopy(preprocess_config)
if isinstance(peak_threshold, (list, tuple)):
# A per-stage [stage 1, stage 2] threshold: stage 1 is the detector.
detector_kwargs["peak_threshold"] = peak_threshold[0]
detector = _build_bbox(
path,
model_type=model_type,
max_instances=max_instances,
nms_iou=bbox_nms_iou,
**detector_kwargs,
)
if not stage_2:
if "centroid" in model_types:
raise ValueError(
"Got a centroid model and a bounding-box detector but no "
"top-down model for them to feed; pass one first stage plus a "
"centered_instance / multi_class_topdown / "
"centered_instance_segmentation model, or the detector alone."
)
assets = detector
else:
if "centroid" in model_types:
raise ValueError(
"Got both a centroid model and a bounding-box detector: pass one "
"of them as the first stage of the top-down pipeline."
)
if len(stage_2) > 1:
raise ValueError(
f"A bounding-box detector feeds one top-down model; got {stage_2}."
)
stage_2_path = model_paths[model_types.index(stage_2[0])]
builder = {
"centered_instance": _build_topdown,
"multi_class_topdown": _build_topdown_multiclass,
}.get(stage_2[0])
if builder is not None:
assets = builder(
None,
stage_2_path,
max_instances=max_instances,
anchor_part=anchor_part,
**common_kwargs,
)
else:
assets = _build_topdown_segmentation(
None,
stage_2_path,
max_instances=max_instances,
anchor_part=anchor_part,
fg_threshold=fg_threshold,
mask_output=mask_output,
polygon_epsilon=polygon_epsilon,
**common_kwargs,
)
assets.detector_assets = detector
elif "embedding" in model_types:
# Appearance-embedding (re-ID). Checked BEFORE the topdown family block:
# a centroid + embedding pair has "centroid" in model_types, so it would
# otherwise enter that block and silently drop the embedding dir.
# - centroid + embedding -> compose centroid -> crop -> embed on raw
# video (TopDownEmbeddingLayer).
# - embedding alone -> single-stage, mask-driven case (EmbeddingLayer).
emb_path = model_paths[model_types.index("embedding")]
centroid_path = (
model_paths[model_types.index("centroid")]
if "centroid" in model_types
else None
)
if centroid_path is not None:
assets = _build_topdown_embedding(
centroid_path,
emb_path,
max_instances=max_instances,
anchor_part=anchor_part,
**common_kwargs,
)
else:
assets = _build_embedding(emb_path, **common_kwargs)
elif "centered_instance_segmentation" in model_types:
# Top-down (crop-centered) instance segmentation. MUST be checked BEFORE
# the topdown family block below: a centroid + centered_instance_segmentation
# pair has "centroid" in model_types, so it would otherwise enter that block
# and — finding neither "centered_instance" nor "multi_class_topdown" — fall
# to the centroid-only else, SILENTLY DROPPING the seg dir (verify/v2 #1).
seg_path = model_paths[model_types.index("centered_instance_segmentation")]
centroid_path = (
model_paths[model_types.index("centroid")]
if "centroid" in model_types
else None
)
assets = _build_topdown_segmentation(
centroid_path,
seg_path,
max_instances=max_instances,
anchor_part=anchor_part,
fg_threshold=fg_threshold,
mask_output=mask_output,
polygon_epsilon=polygon_epsilon,
**common_kwargs,
)
elif (
"centroid" in model_types
or "centered_instance" in model_types
or "multi_class_topdown" in model_types
):
centroid_path = None
confmap_path = None
if "centroid" in model_types:
centroid_path = model_paths[model_types.index("centroid")]
if "centered_instance" in model_types:
confmap_path = model_paths[model_types.index("centered_instance")]
assets = _build_topdown(
centroid_path,
confmap_path,
max_instances=max_instances,
anchor_part=anchor_part,
**common_kwargs,
)
elif "multi_class_topdown" in model_types:
confmap_path = model_paths[model_types.index("multi_class_topdown")]
assets = _build_topdown_multiclass(
centroid_path,
confmap_path,
max_instances=max_instances,
anchor_part=anchor_part,
**common_kwargs,
)
else:
# centroid-only: still goes through _build_topdown with confmap=None
assets = _build_topdown(
centroid_path,
None,
max_instances=max_instances,
anchor_part=anchor_part,
**common_kwargs,
)
else:
raise ValueError(
f"Could not create inference assets from model paths:\n{model_paths}\n"
f"Detected types: {model_types}"
)
_check_crop_stage(assets, model_types)
return assets, model_types