RegressionTracker

class lightning_pose.models.regression_tracker.RegressionTracker(num_keypoints: int, 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', pretrained: bool = True, torch_seed: int = 123, lr_scheduler: str = 'multisteplr', lr_scheduler_params: DictConfig | dict | None = None, **kwargs: Any)[source]

Bases: BaseSupervisedTracker

Base model that produces (x, y) predictions of keypoints from images.

Methods Summary

`forward`(images)	Forward pass through the network.
`get_loss_inputs_labeled`(batch_dict)	Return predicted coordinates for a batch of data.
`predict_step`(batch_dict, batch_idx, **kwargs)	Predict keypoints for a batch of video frames.

Methods Documentation

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

get_loss_inputs_labeled(batch_dict: BaseLabeledBatchDict) → dict[source]: Return predicted coordinates for a batch of data.

predict_step(batch_dict: BaseLabeledBatchDict | UnlabeledBatchDict, batch_idx: int, **kwargs: Any) → Tuple[Tensor, Tensor][source]

Predict 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)