garmentiq.landmark.detection.model_definition

  1"""The HRNet pose model used for landmark detection.
  2
  3`PoseHighResolutionNet` maintains high-resolution feature maps throughout the network
  4rather than recovering resolution at the end, which is what makes it accurate for
  5precise keypoint localisation.
  6"""
  7import os
  8import logging
  9import torch
 10import torch.nn as nn
 11import torch.nn.functional as F
 12
 13
 14BN_MOMENTUM = 0.1
 15logger = logging.getLogger(__name__)
 16
 17
 18def conv3x3(in_planes, out_planes, stride=1):
 19    """
 20    Creates a 3x3 convolutional layer with padding.
 21
 22    Args:
 23        in_planes (int): Number of input channels.
 24        out_planes (int): Number of output channels.
 25        stride (int, optional): Stride of the convolution. Defaults to 1.
 26
 27    Returns:
 28        nn.Conv2d: 3x3 convolution layer with specified parameters.
 29    """
 30    return nn.Conv2d(
 31        in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False
 32    )
 33
 34
 35class BasicBlock(nn.Module):
 36    """
 37    Basic residual block with two 3x3 convolutional layers.
 38
 39    Attributes:
 40        expansion (int): Expansion factor for output channels (always 1 for BasicBlock).
 41        conv1 (nn.Conv2d): First convolutional layer.
 42        bn1 (nn.BatchNorm2d): Batch normalization after first conv.
 43        conv2 (nn.Conv2d): Second convolutional layer.
 44        bn2 (nn.BatchNorm2d): Batch normalization after second conv.
 45        downsample (nn.Module or None): Optional downsampling layer for residual connection.
 46        stride (int): Stride of the first convolution.
 47    """
 48
 49    expansion = 1
 50
 51    def __init__(self, inplanes, planes, stride=1, downsample=None):
 52        """
 53        Initializes BasicBlock.
 54
 55        Args:
 56            inplanes (int): Number of input channels.
 57            planes (int): Number of output channels.
 58            stride (int, optional): Stride of the first convolution. Defaults to 1.
 59            downsample (nn.Module or None, optional): Downsampling layer for residual. Defaults to None.
 60        """
 61        super(BasicBlock, self).__init__()
 62        self.conv1 = conv3x3(inplanes, planes, stride)
 63        self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
 64        self.conv2 = conv3x3(planes, planes)
 65        self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
 66        self.downsample = downsample
 67        self.stride = stride
 68
 69    def forward(self, x):
 70        """
 71        Forward pass of the BasicBlock.
 72
 73        Args:
 74            x (torch.Tensor): Input tensor.
 75
 76        Returns:
 77            torch.Tensor: Output tensor after residual addition and activation.
 78        """
 79        residual = x
 80
 81        out = self.conv1(x)
 82        out = self.bn1(out)
 83        out = F.relu(out, inplace=True)
 84
 85        out = self.conv2(out)
 86        out = self.bn2(out)
 87
 88        if self.downsample is not None:
 89            residual = self.downsample(x)
 90
 91        out += residual
 92        out = F.relu(out, inplace=True)
 93
 94        return out
 95
 96
 97class Bottleneck(nn.Module):
 98    """
 99    Bottleneck residual block with 1x1, 3x3, and 1x1 convolutions.
100
101    Attributes:
102        expansion (int): Expansion factor for output channels (usually 4 for Bottleneck).
103        conv1 (nn.Conv2d): 1x1 convolution reducing channels.
104        bn1 (nn.BatchNorm2d): Batch normalization after conv1.
105        conv2 (nn.Conv2d): 3x3 convolution.
106        bn2 (nn.BatchNorm2d): Batch normalization after conv2.
107        conv3 (nn.Conv2d): 1x1 convolution expanding channels.
108        bn3 (nn.BatchNorm2d): Batch normalization after conv3.
109        downsample (nn.Module or None): Optional downsampling layer for residual connection.
110        stride (int): Stride for the 3x3 convolution.
111    """
112
113    expansion = 4
114
115    def __init__(self, inplanes, planes, stride=1, downsample=None):
116        """
117        Initializes Bottleneck block.
118
119        Args:
120            inplanes (int): Number of input channels.
121            planes (int): Number of output channels before expansion.
122            stride (int, optional): Stride for the 3x3 convolution. Defaults to 1.
123            downsample (nn.Module or None, optional): Downsampling layer for residual. Defaults to None.
124        """
125        super(Bottleneck, self).__init__()
126        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
127        self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
128        self.conv2 = nn.Conv2d(
129            planes, planes, kernel_size=3, stride=stride, padding=1, bias=False
130        )
131        self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
132        self.conv3 = nn.Conv2d(
133            planes, planes * self.expansion, kernel_size=1, bias=False
134        )
135        self.bn3 = nn.BatchNorm2d(planes * self.expansion, momentum=BN_MOMENTUM)
136        self.downsample = downsample
137        self.stride = stride
138
139    def forward(self, x):
140        """
141        Forward pass of the Bottleneck block.
142
143        Args:
144            x (torch.Tensor): Input tensor.
145
146        Returns:
147            torch.Tensor: Output tensor after residual addition and activation.
148        """
149        residual = x
150
151        out = self.conv1(x)
152        out = self.bn1(out)
153        out = F.relu(out, inplace=True)
154
155        out = self.conv2(out)
156        out = self.bn2(out)
157        out = F.relu(out, inplace=True)
158
159        out = self.conv3(out)
160        out = self.bn3(out)
161
162        if self.downsample is not None:
163            residual = self.downsample(x)
164
165        out += residual
166        out = F.relu(out, inplace=True)
167
168        return out
169
170
171class HighResolutionModule(nn.Module):
172    """
173    HighResolutionModule maintains high-resolution representations through multi-branch architecture and fusion.
174
175    This module consists of several parallel branches with residual blocks, and fuse layers to combine
176    features from different resolutions.
177
178    Attributes:
179        num_branches (int): Number of parallel branches.
180        blocks (nn.Module): Residual block class (BasicBlock or Bottleneck).
181        num_blocks (list[int]): Number of residual blocks per branch.
182        num_inchannels (list[int]): Number of input channels for each branch.
183        num_channels (list[int]): Number of channels per branch before expansion.
184        fuse_method (str): Method to fuse multi-branch outputs ('SUM' supported).
185        multi_scale_output (bool): Whether to output multi-scale features.
186        branches (nn.ModuleList): The parallel branches.
187        fuse_layers (nn.ModuleList or None): Layers that fuse features from branches.
188    """
189
190    def __init__(
191        self,
192        num_branches,
193        blocks,
194        num_blocks,
195        num_inchannels,
196        num_channels,
197        fuse_method,
198        multi_scale_output=True,
199    ):
200        """
201        Initializes HighResolutionModule.
202
203        Args:
204            num_branches (int): Number of parallel branches.
205            blocks (nn.Module): Residual block class (BasicBlock or Bottleneck).
206            num_blocks (list[int]): Number of residual blocks per branch.
207            num_inchannels (list[int]): Number of input channels for each branch.
208            num_channels (list[int]): Number of channels per branch before expansion.
209            fuse_method (str): Method to fuse multi-branch outputs.
210            multi_scale_output (bool, optional): Output multi-scale features or not. Defaults to True.
211
212        Raises:
213            ValueError: If lengths of inputs do not match num_branches.
214        """
215        super(HighResolutionModule, self).__init__()
216        self._check_branches(
217            num_branches, blocks, num_blocks, num_inchannels, num_channels
218        )
219
220        self.num_inchannels = num_inchannels
221        self.fuse_method = fuse_method
222        self.num_branches = num_branches
223
224        self.multi_scale_output = multi_scale_output
225
226        self.branches = self._make_branches(
227            num_branches, blocks, num_blocks, num_channels
228        )
229        self.fuse_layers = self._make_fuse_layers()
230
231    def _check_branches(
232        self, num_branches, blocks, num_blocks, num_inchannels, num_channels
233    ):
234        """
235        Validates that lengths of num_blocks, num_inchannels, and num_channels match num_branches.
236
237        Args:
238            num_branches (int): Number of branches.
239            blocks (nn.Module): Block type.
240            num_blocks (list[int]): Number of blocks per branch.
241            num_inchannels (list[int]): Number of input channels per branch.
242            num_channels (list[int]): Number of channels per branch.
243
244        Raises:
245            ValueError: If any length mismatch occurs.
246        """
247        if num_branches != len(num_blocks):
248            error_msg = "NUM_BRANCHES({}) <> NUM_BLOCKS({})".format(
249                num_branches, len(num_blocks)
250            )
251            logger.error(error_msg)
252            raise ValueError(error_msg)
253
254        if num_branches != len(num_channels):
255            error_msg = "NUM_BRANCHES({}) <> NUM_CHANNELS({})".format(
256                num_branches, len(num_channels)
257            )
258            logger.error(error_msg)
259            raise ValueError(error_msg)
260
261        if num_branches != len(num_inchannels):
262            error_msg = "NUM_BRANCHES({}) <> NUM_INCHANNELS({})".format(
263                num_branches, len(num_inchannels)
264            )
265            logger.error(error_msg)
266            raise ValueError(error_msg)
267
268    def _make_one_branch(self, branch_index, block, num_blocks, num_channels, stride=1):
269        """
270        Constructs one branch of the module consisting of sequential residual blocks.
271
272        Args:
273            branch_index (int): Index of the branch.
274            block (nn.Module): Residual block class.
275            num_blocks (int): Number of residual blocks in this branch.
276            num_channels (list[int]): Number of channels per branch.
277            stride (int, optional): Stride of the first block. Defaults to 1.
278
279        Returns:
280            nn.Sequential: Sequential container of residual blocks.
281        """
282        downsample = None
283        if (
284            stride != 1
285            or self.num_inchannels[branch_index]
286            != num_channels[branch_index] * block.expansion
287        ):
288            downsample = nn.Sequential(
289                nn.Conv2d(
290                    self.num_inchannels[branch_index],
291                    num_channels[branch_index] * block.expansion,
292                    kernel_size=1,
293                    stride=stride,
294                    bias=False,
295                ),
296                nn.BatchNorm2d(
297                    num_channels[branch_index] * block.expansion, momentum=BN_MOMENTUM
298                ),
299            )
300
301        layers = []
302        layers.append(
303            block(
304                self.num_inchannels[branch_index],
305                num_channels[branch_index],
306                stride,
307                downsample,
308            )
309        )
310        self.num_inchannels[branch_index] = num_channels[branch_index] * block.expansion
311        for i in range(1, num_blocks[branch_index]):
312            layers.append(
313                block(self.num_inchannels[branch_index], num_channels[branch_index])
314            )
315
316        return nn.Sequential(*layers)
317
318    def _make_branches(self, num_branches, block, num_blocks, num_channels):
319        """
320        Constructs all branches for the module.
321
322        Args:
323            num_branches (int): Number of branches.
324            block (nn.Module): Residual block class.
325            num_blocks (list[int]): Number of blocks per branch.
326            num_channels (list[int]): Number of channels per branch.
327
328        Returns:
329            nn.ModuleList: List of branch modules.
330        """
331        branches = []
332
333        for i in range(num_branches):
334            branches.append(self._make_one_branch(i, block, num_blocks, num_channels))
335
336        return nn.ModuleList(branches)
337
338    def _make_fuse_layers(self):
339        """
340        Constructs layers to fuse multi-resolution branch outputs.
341
342        Returns:
343            nn.ModuleList or None: Fuse layers or None if single branch.
344        """
345        if self.num_branches == 1:
346            return None
347
348        num_branches = self.num_branches
349        num_inchannels = self.num_inchannels
350        fuse_layers = []
351        for i in range(num_branches if self.multi_scale_output else 1):
352            fuse_layer = []
353            for j in range(num_branches):
354                if j > i:
355                    fuse_layer.append(
356                        nn.Sequential(
357                            nn.Conv2d(
358                                num_inchannels[j],
359                                num_inchannels[i],
360                                1,
361                                1,
362                                0,
363                                bias=False,
364                            ),
365                            nn.BatchNorm2d(num_inchannels[i]),
366                            nn.Upsample(scale_factor=2 ** (j - i), mode="nearest"),
367                        )
368                    )
369                elif j == i:
370                    fuse_layer.append(None)
371                else:
372                    conv3x3s = []
373                    for k in range(i - j):
374                        if k == i - j - 1:
375                            num_outchannels_conv3x3 = num_inchannels[i]
376                            conv3x3s.append(
377                                nn.Sequential(
378                                    nn.Conv2d(
379                                        num_inchannels[j],
380                                        num_outchannels_conv3x3,
381                                        3,
382                                        2,
383                                        1,
384                                        bias=False,
385                                    ),
386                                    nn.BatchNorm2d(num_outchannels_conv3x3),
387                                )
388                            )
389                        else:
390                            num_outchannels_conv3x3 = num_inchannels[j]
391                            conv3x3s.append(
392                                nn.Sequential(
393                                    nn.Conv2d(
394                                        num_inchannels[j],
395                                        num_outchannels_conv3x3,
396                                        3,
397                                        2,
398                                        1,
399                                        bias=False,
400                                    ),
401                                    nn.BatchNorm2d(num_outchannels_conv3x3),
402                                    nn.ReLU(True),
403                                )
404                            )
405                    fuse_layer.append(nn.Sequential(*conv3x3s))
406            fuse_layers.append(nn.ModuleList(fuse_layer))
407
408        return nn.ModuleList(fuse_layers)
409
410    def get_num_inchannels(self):
411        """
412        Returns the number of input channels for each branch after block expansion.
413
414        Returns:
415            list[int]: Number of input channels per branch.
416        """
417        return self.num_inchannels
418
419    def forward(self, x):
420        """
421        Forward pass through the HighResolutionModule.
422
423        Args:
424            x (list[torch.Tensor]): List of input tensors for each branch.
425
426        Returns:
427            list[torch.Tensor]: List of output tensors after multi-branch fusion.
428        """
429        if self.num_branches == 1:
430            return [self.branches[0](x[0])]
431
432        for i in range(self.num_branches):
433            x[i] = self.branches[i](x[i])
434
435        x_fuse = []
436
437        for i in range(len(self.fuse_layers)):
438            y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
439            for j in range(1, self.num_branches):
440                if i == j:
441                    y = y + x[j]
442                else:
443                    y = y + self.fuse_layers[i][j](x[j])
444            x_fuse.append(F.relu(y, inplace=True))
445
446        return x_fuse
447
448
449blocks_dict = {"BASIC": BasicBlock, "BOTTLENECK": Bottleneck}
450
451
452class PoseHighResolutionNet(nn.Module):
453    """
454    High-Resolution Network (HRNet) tailored for garment landmark detection.
455
456    The network maintains high-resolution representations through multiple stages and branches,
457    fusing multi-scale features and finally predicting heatmaps or coordinates for landmarks.
458
459    Attributes:
460        inplanes (int): Initial number of input channels.
461        conv1 (nn.Conv2d): Initial 3x3 convolution.
462        bn1 (nn.BatchNorm2d): Batch normalization after conv1.
463        conv2 (nn.Conv2d): Second 3x3 convolution.
464        bn2 (nn.BatchNorm2d): Batch normalization after conv2.
465        relu (nn.ReLU): ReLU activation.
466        stage1_cfg (dict): Configuration for stage1.
467        stage2_cfg (dict): Configuration for stage2.
468        stage3_cfg (dict): Configuration for stage3.
469        stage4_cfg (dict): Configuration for stage4.
470        transition1 (nn.ModuleList): Transition layers between stages.
471        stage2 (HighResolutionModule): Stage 2 module.
472        transition2 (nn.ModuleList): Transition layers between stages.
473        stage3 (HighResolutionModule): Stage 3 module.
474        transition3 (nn.ModuleList): Transition layers between stages.
475        stage4 (HighResolutionModule): Stage 4 module.
476        final_layer (nn.Conv2d): Final convolution layer to output predictions.
477        target_type (str): Output format type ('gaussian' or 'coordinate').
478    """
479
480    def __init__(self, **kwargs):
481        """
482        Initializes PoseHighResolutionNet with default HRNet configurations.
483
484        Args:
485            target_type (str, optional): Type of target output. Either "gaussian" or "coordinate". Defaults to "gaussian".
486        """
487        self.inplanes = 64
488        # Hardcoded values from YAML MODEL.EXTRA
489        extra = {
490            "PRETRAINED_LAYERS": [
491                "conv1",
492                "bn1",
493                "conv2",
494                "bn2",
495                "layer1",
496                "transition1",
497                "stage2",
498                "transition2",
499                "stage3",
500                "transition3",
501                "stage4",
502            ],
503            "FINAL_CONV_KERNEL": 1,
504            "STAGE2": {
505                "NUM_MODULES": 1,
506                "NUM_BRANCHES": 2,
507                "BLOCK": "BASIC",
508                "NUM_BLOCKS": [4, 4],
509                "NUM_CHANNELS": [48, 96],
510                "FUSE_METHOD": "SUM",
511            },
512            "STAGE3": {
513                "NUM_MODULES": 4,
514                "NUM_BRANCHES": 3,
515                "BLOCK": "BASIC",
516                "NUM_BLOCKS": [4, 4, 4],
517                "NUM_CHANNELS": [48, 96, 192],
518                "FUSE_METHOD": "SUM",
519            },
520            "STAGE4": {
521                "NUM_MODULES": 3,
522                "NUM_BRANCHES": 4,
523                "BLOCK": "BASIC",
524                "NUM_BLOCKS": [4, 4, 4, 4],
525                "NUM_CHANNELS": [48, 96, 192, 384],
526                "FUSE_METHOD": "SUM",
527            },
528        }
529
530        self.model_name = "pose_hrnet"
531        self.target_type = "gaussian"
532
533        super(PoseHighResolutionNet, self).__init__()
534
535        # stem net
536        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False)
537        self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
538        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1, bias=False)
539        self.bn2 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
540        self.layer1 = self._make_layer(Bottleneck, 64, 4)
541
542        # Stage2
543        self.stage2_cfg = extra["STAGE2"]
544        num_channels = self.stage2_cfg["NUM_CHANNELS"]
545        block = blocks_dict[self.stage2_cfg["BLOCK"]]
546        num_channels = [
547            num_channels[i] * block.expansion for i in range(len(num_channels))
548        ]
549        self.transition1 = self._make_transition_layer([256], num_channels)
550        self.stage2, pre_stage_channels = self._make_stage(
551            self.stage2_cfg, num_channels
552        )
553
554        # Stage3
555        self.stage3_cfg = extra["STAGE3"]
556        num_channels = self.stage3_cfg["NUM_CHANNELS"]
557        block = blocks_dict[self.stage3_cfg["BLOCK"]]
558        num_channels = [
559            num_channels[i] * block.expansion for i in range(len(num_channels))
560        ]
561        self.transition2 = self._make_transition_layer(pre_stage_channels, num_channels)
562        self.stage3, pre_stage_channels = self._make_stage(
563            self.stage3_cfg, num_channels
564        )
565
566        # Stage4
567        self.stage4_cfg = extra["STAGE4"]
568        num_channels = self.stage4_cfg["NUM_CHANNELS"]
569        block = blocks_dict[self.stage4_cfg["BLOCK"]]
570        num_channels = [
571            num_channels[i] * block.expansion for i in range(len(num_channels))
572        ]
573        self.transition3 = self._make_transition_layer(pre_stage_channels, num_channels)
574        self.stage4, pre_stage_channels = self._make_stage(
575            self.stage4_cfg, num_channels, multi_scale_output=False
576        )
577
578        # Final layer
579        self.final_layer = nn.Conv2d(
580            in_channels=pre_stage_channels[0],
581            out_channels=294,  # from MODEL.NUM_JOINTS
582            kernel_size=extra["FINAL_CONV_KERNEL"],
583            stride=1,
584            padding=1 if extra["FINAL_CONV_KERNEL"] == 3 else 0,
585        )
586
587        self.pretrained_layers = extra["PRETRAINED_LAYERS"]
588
589    def _make_transition_layer(self, num_channels_pre_layer, num_channels_cur_layer):
590        """
591        Creates transition layers to match the number of channels between stages.
592
593        Args:
594            num_channels_pre_layer (list[int]): Channels from previous stage.
595            num_channels_cur_layer (list[int]): Channels for current stage.
596
597        Returns:
598            nn.ModuleList: List of transition layers.
599        """
600        num_branches_cur = len(num_channels_cur_layer)
601        num_branches_pre = len(num_channels_pre_layer)
602
603        transition_layers = []
604        for i in range(num_branches_cur):
605            if i < num_branches_pre:
606                if num_channels_cur_layer[i] != num_channels_pre_layer[i]:
607                    transition_layers.append(
608                        nn.Sequential(
609                            nn.Conv2d(
610                                num_channels_pre_layer[i],
611                                num_channels_cur_layer[i],
612                                3,
613                                1,
614                                1,
615                                bias=False,
616                            ),
617                            nn.BatchNorm2d(num_channels_cur_layer[i]),
618                            nn.ReLU(inplace=True),
619                        )
620                    )
621                else:
622                    transition_layers.append(None)
623            else:
624                conv3x3s = []
625                for j in range(i + 1 - num_branches_pre):
626                    inchannels = num_channels_pre_layer[-1]
627                    outchannels = (
628                        num_channels_cur_layer[i]
629                        if j == i - num_branches_pre
630                        else inchannels
631                    )
632                    conv3x3s.append(
633                        nn.Sequential(
634                            nn.Conv2d(inchannels, outchannels, 3, 2, 1, bias=False),
635                            nn.BatchNorm2d(outchannels),
636                            nn.ReLU(inplace=True),
637                        )
638                    )
639                transition_layers.append(nn.Sequential(*conv3x3s))
640
641        return nn.ModuleList(transition_layers)
642
643    def _make_layer(self, block, planes, blocks, stride=1):
644        """
645        Creates a layer composed of sequential residual blocks.
646
647        Args:
648            block (nn.Module): Residual block type (BasicBlock or Bottleneck).
649            planes (int): Number of output channels.
650            blocks (int): Number of blocks in this layer.
651            stride (int, optional): Stride for the first block. Defaults to 1.
652
653        Returns:
654            nn.Sequential: Sequential container of residual blocks.
655        """
656        downsample = None
657        if stride != 1 or self.inplanes != planes * block.expansion:
658            downsample = nn.Sequential(
659                nn.Conv2d(
660                    self.inplanes,
661                    planes * block.expansion,
662                    kernel_size=1,
663                    stride=stride,
664                    bias=False,
665                ),
666                nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM),
667            )
668
669        layers = []
670        layers.append(block(self.inplanes, planes, stride, downsample))
671        self.inplanes = planes * block.expansion
672        for i in range(1, blocks):
673            layers.append(block(self.inplanes, planes))
674
675        return nn.Sequential(*layers)
676
677    def _make_stage(self, layer_config, num_inchannels, multi_scale_output=True):
678        """
679        Constructs a stage consisting of one or more HighResolutionModules.
680
681        Args:
682            layer_config (dict): Configuration dictionary for the stage.
683            num_inchannels (list[int]): Number of input channels for each branch.
684            multi_scale_output (bool, optional): Output multi-scale features or not. Defaults to True.
685
686        Returns:
687            tuple:
688                nn.Sequential: Stage module.
689                list[int]: Number of output channels for each branch.
690        """
691        num_modules = layer_config["NUM_MODULES"]
692        num_branches = layer_config["NUM_BRANCHES"]
693        num_blocks = layer_config["NUM_BLOCKS"]
694        num_channels = layer_config["NUM_CHANNELS"]
695        block = blocks_dict[layer_config["BLOCK"]]
696        fuse_method = layer_config["FUSE_METHOD"]
697
698        modules = []
699        for i in range(num_modules):
700            # multi_scale_output is only used last module
701            if not multi_scale_output and i == num_modules - 1:
702                reset_multi_scale_output = False
703            else:
704                reset_multi_scale_output = True
705
706            modules.append(
707                HighResolutionModule(
708                    num_branches,
709                    block,
710                    num_blocks,
711                    num_inchannels,
712                    num_channels,
713                    fuse_method,
714                    reset_multi_scale_output,
715                )
716            )
717            num_inchannels = modules[-1].get_num_inchannels()
718
719        return nn.Sequential(*modules), num_inchannels
720
721    def forward(self, x):
722        """
723        Forward pass of the PoseHighResolutionNet.
724
725        Args:
726            x (torch.Tensor): Input tensor of shape (batch_size, 3, H, W).
727
728        Returns:
729            torch.Tensor: Output heatmaps or coordinates for landmarks.
730        """
731        x = self.conv1(x)
732        x = self.bn1(x)
733        x = F.relu(x, inplace=True)
734        x = self.conv2(x)
735        x = self.bn2(x)
736        x = F.relu(x, inplace=True)
737        x = self.layer1(x)
738
739        x_list = []
740        for i in range(self.stage2_cfg["NUM_BRANCHES"]):
741            if self.transition1[i] is not None:
742                x_list.append(self.transition1[i](x))
743            else:
744                x_list.append(x)
745        y_list = self.stage2(x_list)
746
747        x_list = []
748        for i in range(self.stage3_cfg["NUM_BRANCHES"]):
749            if self.transition2[i] is not None:
750                x_list.append(self.transition2[i](y_list[-1]))
751            else:
752                x_list.append(y_list[i])
753        y_list = self.stage3(x_list)
754
755        x_list = []
756        for i in range(self.stage4_cfg["NUM_BRANCHES"]):
757            if self.transition3[i] is not None:
758                x_list.append(self.transition3[i](y_list[-1]))
759            else:
760                x_list.append(y_list[i])
761        y_list = self.stage4(x_list)
762
763        if self.model_name == "pose_hrnet" or "pose_metric_gcn":
764            x = self.final_layer(y_list[0])
765        else:
766            x = y_list[0]
767
768        if self.target_type == "gaussian":
769            return x
770
771        elif self.target_type == "coordinate":
772            B, C, H, W = x.shape
773
774            """B - cal x,y seperately"""
775            h = F.softmax(x.view(B, C, H * W) * 1, dim=2)
776            h = h.view(B, C, H, W)
777            hx = h.sum(dim=2)  # (B, C, W)
778            px = (hx * (torch.arange(W, device=h.device).float().view(1, 1, W))).sum(
779                2, keepdim=True
780            )
781            hy = h.sum(dim=3)  # (B, C, H)
782            py = (hy * (torch.arange(H, device=h.device).float().view(1, 1, H))).sum(
783                2, keepdim=True
784            )
785            x = torch.cat([px, py], dim=2)
786            return h, x
787        else:
788            raise NotImplementedError(f"{self.target_type} is unknown.")