LinearRegressionHead

class lightning_pose.models.heads.LinearRegressionHead

Bases: Module

Linear regression head that converts 2D feature maps to a vector of (x, y) coordinates.

Methods Summary

forward(features)

Map feature maps to keypoint coordinate predictions.

Methods Documentation

forward(features: Float[Tensor, 'batch features height width']) → Float[Tensor, 'batch coordinates']

Map feature maps to keypoint coordinate predictions.

Parameters:

features – feature tensor of shape (batch, features, height, width); spatial dimensions are collapsed before the linear layer.

Returns:

Predicted coordinates of shape (batch, num_targets).

__init__(in_channels: int, num_targets: int) → None

Initialize LinearRegressionHead.

Parameters:

• in_channels – number of input feature channels.

• num_targets – number of output coordinate values (2 * num_keypoints).

__new__(**kwargs)