HeatmapTracker

class lightning_pose.models.heatmap_tracker.HeatmapTracker(num_keypoints: int, num_targets: int | None = None, loss_factory: LossFactory | None = None, backbone: Literal['resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152', 'resnet50_contrastive', 'resnet50_animal_apose', 'resnet50_animal_ap10k', 'resnet50_human_jhmdb', 'resnet50_human_res_rle', 'resnet50_human_top_res', 'resnet50_human_hand', 'efficientnet_b0', 'efficientnet_b1', 'efficientnet_b2', 'vit_b_sam'] = 'resnet50', downsample_factor: typing_extensions.Literal[1, 2, 3] = 2, pretrained: bool = True, output_shape: tuple | None = None, torch_seed: int = 123, lr_scheduler: str = 'multisteplr', lr_scheduler_params: DictConfig | dict | None = None, **kwargs: Any)[source]

Bases: BaseSupervisedTracker

Base model that produces heatmaps of keypoints from images.

Attributes Summary

num_filters_for_upsampling

Methods Summary

`create_double_upsampling_layer`(in_channels, ...)	Perform ConvTranspose2d to double the output shape.
`forward`(images)	Forward pass through the network.
`get_loss_inputs_labeled`(batch_dict)	Return predicted heatmaps and their softmaxes (estimated keypoints).
`heatmaps_from_representations`(representations)	Upsample representations to get final heatmaps.
`initialize_upsampling_layers`()	Intialize the Conv2DTranspose upsampling layers.
`make_upsampling_layers`()
`predict_step`(batch_dict, batch_idx[, ...])	Predict heatmaps and keypoints for a batch of video frames.
`run_hard_argmax`(heatmaps)	Use hard argmax on heatmaps.
`run_subpixelmaxima`(heatmaps)	Use soft argmax on heatmaps.

Attributes Documentation

num_filters_for_upsampling

Methods Documentation

static create_double_upsampling_layer(in_channels: int, out_channels: int) → ConvTranspose2d[source]: Perform ConvTranspose2d to double the output shape.

forward(images: Tensor[Tensor] | Tensor[Tensor]) → Tensor[Tensor][source]: Forward pass through the network.

get_loss_inputs_labeled(batch_dict: HeatmapLabeledBatchDict | MultiviewHeatmapLabeledBatchDict) → dict[source]: Return predicted heatmaps and their softmaxes (estimated keypoints).

heatmaps_from_representations(representations: Tensor[Tensor]) → Tensor[Tensor][source]: Upsample representations to get final heatmaps.

initialize_upsampling_layers() → None[source]: Intialize the Conv2DTranspose upsampling layers.

make_upsampling_layers() → Sequential[source]

predict_step(batch_dict: HeatmapLabeledBatchDict | MultiviewHeatmapLabeledBatchDict | UnlabeledBatchDict, batch_idx: int, return_heatmaps: bool | None = False) → Tuple[Tensor, Tensor] | Tuple[Tensor, Tensor, Tensor][source]

Predict heatmaps and keypoints for a batch of video frames.

Assuming a DALI video loader is passed in > trainer = Trainer(devices=8, accelerator=”gpu”) > predictions = trainer.predict(model, data_loader)

run_hard_argmax(heatmaps: Tensor[Tensor]) → Tuple[Tensor[Tensor], Tensor[Tensor]][source]

Use hard argmax on heatmaps.

Parameters:

heatmaps – output of upsampling layers

Returns:

tuple

hard argmax of shape (batch, num_targets)
confidences of shape (batch, num_keypoints)

run_subpixelmaxima(heatmaps: Tensor[Tensor]) → Tuple[Tensor[Tensor], Tensor[Tensor]][source]

Use soft argmax on heatmaps.

Parameters:

heatmaps – output of upsampling layers

Returns:

tuple

soft argmax of shape (batch, num_targets)
confidences of shape (batch, num_keypoints)