Skip to content

deeplabcut.pose_estimation_pytorch.models.backbones.cond_prenet

Classes:

Name Description
CondPreNet

Wrapper module that adds a conditional preNet before any backbone.

CondPreNet

Bases: BaseBackbone

Wrapper module that adds a conditional preNet before any backbone.

This allows to process image and condition features and prepare them for the main backbone.

Methods:

Name Description
__init__

Initialize the PreNetWrapper.
forward

Forward pass through the conditional preNet + backbone.
init_weights

Initialize PreNet weights from a Normal distribution.
Source code in deeplabcut/pose_estimation_pytorch/models/backbones/cond_prenet.py 
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
@BACKBONES.register_module
class CondPreNet(BaseBackbone):
    """Wrapper module that adds a conditional preNet before any backbone.

    This allows to process image and condition features and prepare them for the main
    backbone.
    """

    def __init__(
        self,
        kpt_encoder: dict | BaseKeypointEncoder,
        backbone: dict | BaseBackbone,
        img_size: tuple[int, int] = (256, 256),
        **kwargs,
    ):
        """Initialize the PreNetWrapper.

        Args:
            backbone: The backbone model to wrap.
            img_size: The (height, width) of the input images.
        """
        pretrained = kwargs.pop("pretrained", False)
        if not isinstance(backbone, BaseBackbone):
            backbone["pretrained"] = pretrained
            backbone = BACKBONES.build(backbone)

        super().__init__(stride=backbone.stride, **kwargs)

        if not isinstance(kpt_encoder, BaseKeypointEncoder):
            if "img_size" not in kpt_encoder:
                kpt_encoder["img_size"] = img_size
            kpt_encoder = KEYPOINT_ENCODERS.build(kpt_encoder)
        self.cond_enc = kpt_encoder

        self.backbone = backbone
        self.rgb_preNet = self._make_preNet(num_inputs=3, num_outputs=3, input_image=True)
        self.cond_preNet = self._make_preNet(num_inputs=self.cond_enc.num_channels, num_outputs=3, input_image=False)

        self.init_weights()

    def _make_preNet(self, num_inputs, num_outputs, input_image=False):
        if not input_image:  # cond
            preNet = nn.Sequential(
                nn.Conv2d(num_inputs, num_outputs, kernel_size=7, stride=1, padding="same"),
                nn.BatchNorm2d(num_outputs),
            )
        else:
            preNet = nn.Sequential(
                nn.Conv2d(num_inputs, 64, kernel_size=3, stride=1, padding="same"),
                nn.BatchNorm2d(64),
                nn.Conv2d(64, num_outputs, kernel_size=7, stride=1, padding="same"),
                nn.BatchNorm2d(num_outputs),
            )
        return preNet

    def forward(self, x: torch.Tensor, cond_kpts: np.ndarray | torch.Tensor) -> torch.Tensor:
        """Forward pass through the conditional preNet + backbone.

        Args:
            x: Input tensor of shape (batch_size, channels, height, width).
            cond_kpts: Conditional keypoints of shape (batch_size, num_joints, 2).

        Returns:
            the feature map
        """
        # create conditional heatmap
        if isinstance(cond_kpts, torch.Tensor):
            cond_kpts = cond_kpts.detach().numpy()

        cond_hm = self.cond_enc(cond_kpts.squeeze(1), x.size()[2:])
        cond_hm = torch.from_numpy(cond_hm).float().to(x.device)
        cond_hm = cond_hm.permute(0, 3, 1, 2)  # (B, C, H, W)

        x0 = self.rgb_preNet(x)
        x1 = self.cond_preNet(cond_hm)
        x = x0 + x1

        return self.backbone(x)

    def init_weights(self):
        """Initialize PreNet weights from a Normal distribution."""
        for prenet in [self.rgb_preNet, self.cond_preNet]:
            for m in prenet.modules():
                if isinstance(m, nn.Conv2d):
                    nn.init.normal_(m.weight, std=0.001)
                    if m.bias is not None:
                        nn.init.constant_(m.bias, 0)
                elif isinstance(m, nn.BatchNorm2d):
                    nn.init.constant_(m.weight, 1)
                    nn.init.constant_(m.bias, 0)

__init__ 

__init__(
    kpt_encoder: dict | BaseKeypointEncoder,
    backbone: dict | BaseBackbone,
    img_size: tuple[int, int] = (256, 256),
    **kwargs
)

Initialize the PreNetWrapper.

Parameters:

Name Type Description Default

backbone

dict | BaseBackbone

The backbone model to wrap.

 required

img_size

tuple[int, int]

The (height, width) of the input images.

 (256, 256)
Source code in deeplabcut/pose_estimation_pytorch/models/backbones/cond_prenet.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
def __init__(
    self,
    kpt_encoder: dict | BaseKeypointEncoder,
    backbone: dict | BaseBackbone,
    img_size: tuple[int, int] = (256, 256),
    **kwargs,
):
    """Initialize the PreNetWrapper.

    Args:
        backbone: The backbone model to wrap.
        img_size: The (height, width) of the input images.
    """
    pretrained = kwargs.pop("pretrained", False)
    if not isinstance(backbone, BaseBackbone):
        backbone["pretrained"] = pretrained
        backbone = BACKBONES.build(backbone)

    super().__init__(stride=backbone.stride, **kwargs)

    if not isinstance(kpt_encoder, BaseKeypointEncoder):
        if "img_size" not in kpt_encoder:
            kpt_encoder["img_size"] = img_size
        kpt_encoder = KEYPOINT_ENCODERS.build(kpt_encoder)
    self.cond_enc = kpt_encoder

    self.backbone = backbone
    self.rgb_preNet = self._make_preNet(num_inputs=3, num_outputs=3, input_image=True)
    self.cond_preNet = self._make_preNet(num_inputs=self.cond_enc.num_channels, num_outputs=3, input_image=False)

    self.init_weights()

forward 

forward(x: Tensor, cond_kpts: ndarray | Tensor) -> torch.Tensor

Forward pass through the conditional preNet + backbone.

Parameters:

Name Type Description Default

x

Tensor

Input tensor of shape (batch_size, channels, height, width).

 required

cond_kpts

ndarray | Tensor

Conditional keypoints of shape (batch_size, num_joints, 2).

 required

Returns:

Type Description
Tensor

the feature map
Source code in deeplabcut/pose_estimation_pytorch/models/backbones/cond_prenet.py 
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def forward(self, x: torch.Tensor, cond_kpts: np.ndarray | torch.Tensor) -> torch.Tensor:
    """Forward pass through the conditional preNet + backbone.

    Args:
        x: Input tensor of shape (batch_size, channels, height, width).
        cond_kpts: Conditional keypoints of shape (batch_size, num_joints, 2).

    Returns:
        the feature map
    """
    # create conditional heatmap
    if isinstance(cond_kpts, torch.Tensor):
        cond_kpts = cond_kpts.detach().numpy()

    cond_hm = self.cond_enc(cond_kpts.squeeze(1), x.size()[2:])
    cond_hm = torch.from_numpy(cond_hm).float().to(x.device)
    cond_hm = cond_hm.permute(0, 3, 1, 2)  # (B, C, H, W)

    x0 = self.rgb_preNet(x)
    x1 = self.cond_preNet(cond_hm)
    x = x0 + x1

    return self.backbone(x)

init_weights 

init_weights()

Initialize PreNet weights from a Normal distribution.

Source code in deeplabcut/pose_estimation_pytorch/models/backbones/cond_prenet.py 
104
105
106
107
108
109
110
111
112
113
114
def init_weights(self):
    """Initialize PreNet weights from a Normal distribution."""
    for prenet in [self.rgb_preNet, self.cond_preNet]:
        for m in prenet.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.normal_(m.weight, std=0.001)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)