deeplabcut.pose_estimation_pytorch.models.target_generators.base

Classes:

Name	Description
BaseGenerator	Generates target maps from ground truth annotations to train models.

BaseGenerator

Bases: ABC, Module

Generates target maps from ground truth annotations to train models.

The outputs of the target generator are used to compute losses for model heads. If the head outputs "heatmap" and "offset" tensors, then the corresponding generator must output target "heatmap" and "offset" tensors. The targets themselves are dictionaries, and passed as keyword-arguments to the criterions. This allows to pass masks to the criterions.

Generally, this means that for each head output (such as "heatmap"), a dict will be generated with a "target" key (for the target heatmap) and optionally a "weights" key (see the WeightedCriterion classes).

Methods:

Name	Description
forward	Generates targets.

Source code in deeplabcut/pose_estimation_pytorch/models/target_generators/base.py

class BaseGenerator(ABC, nn.Module):  # TODO: Should this really be a module?
    """Generates target maps from ground truth annotations to train models.

    The outputs of the target generator are used to compute losses for model heads. If
    the head outputs "heatmap" and "offset" tensors, then the corresponding generator
    must output target "heatmap" and "offset" tensors. The targets themselves are
    dictionaries, and passed as keyword-arguments to the criterions. This allows to pass
    masks to the criterions.

    Generally, this means that for each head output (such as "heatmap"), a dict will be
    generated with a "target" key (for the target heatmap) and optionally a "weights"
    key (see the WeightedCriterion classes).
    """

    def __init__(self, label_keypoint_key: str = "keypoints"):
        super().__init__()
        self.label_keypoint_key = label_keypoint_key

    @abstractmethod
    def forward(
        self, stride: float, outputs: dict[str, torch.Tensor], labels: dict
    ) -> dict[str, dict[str, torch.Tensor]]:
        """Generates targets.

        Args:
            stride: the stride of the model
            outputs: output of a model head
            labels: the labels for the inputs (each tensor should have shape (b, ...))

        Returns:
            a dictionary mapping the heads to the inputs of the criterion
                {
                    "heatmap": {
                        "target": heatmaps,
                        "weights":  heatmap_weights,
                    },
                    "locref": {
                        "target": locref_map,
                        "weights": locref_weights,
                    }
                }
        """

forward abstractmethod

forward(stride: float, outputs: dict[str, Tensor], labels: dict) -> dict[str, dict[str, torch.Tensor]]

Generates targets.

Parameters:

Name	Type	Description	Default
stride	float	the stride of the model	required
outputs	dict[str, Tensor]	output of a model head	required
labels	dict	the labels for the inputs (each tensor should have shape (b, ...))	required

Returns:

Type	Description
dict[str, dict[str, Tensor]]	a dictionary mapping the heads to the inputs of the criterion { "heatmap": { "target": heatmaps, "weights": heatmap_weights, }, "locref": { "target": locref_map, "weights": locref_weights, } }

Source code in deeplabcut/pose_estimation_pytorch/models/target_generators/base.py

@abstractmethod
def forward(
    self, stride: float, outputs: dict[str, torch.Tensor], labels: dict
) -> dict[str, dict[str, torch.Tensor]]:
    """Generates targets.

    Args:
        stride: the stride of the model
        outputs: output of a model head
        labels: the labels for the inputs (each tensor should have shape (b, ...))

    Returns:
        a dictionary mapping the heads to the inputs of the criterion
            {
                "heatmap": {
                    "target": heatmaps,
                    "weights":  heatmap_weights,
                },
                "locref": {
                    "target": locref_map,
                    "weights": locref_weights,
                }
            }
    """