modeling_rtmw.py

from typing import Optional, Tuple, Union, Dict, Sequence
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers.modeling_outputs import ModelOutput
from transformers.modeling_utils import PreTrainedModel
from transformers.utils import logging

from .configuration_rtmw import RTMWConfig


logger = logging.get_logger(__name__)


@dataclass
class PoseOutput(ModelOutput):
    """
    Output type for pose estimation models.
    
    Args:
        keypoints (`torch.FloatTensor` of shape `(batch_size, num_keypoints, 2)`):
            Predicted keypoint coordinates in format [x, y]. The coordinate system
            depends on the `coordinate_mode` passed to `forward()`:
            - ``"model"``        — raw SimCC space (model input resolution, e.g. 288×384 px)
            - ``"image"``        — original image space, scaled via the supplied `bbox`
            - ``"root_relative"`` — root-normalised: origin at mid-hip, unit = half hip-to-hip dist
        scores (`torch.FloatTensor` of shape `(batch_size, num_keypoints)`):
            Predicted keypoint confidence scores in [0, 1].
        coordinate_mode (`str`):
            Which coordinate system `keypoints` is expressed in (mirrors the arg passed to forward).
        loss (`torch.FloatTensor`, *optional*):
            Loss value if training.
        pred_x (`torch.FloatTensor`, *optional*):
            X-axis heatmap predictions from the SimCC representation.
        pred_y (`torch.FloatTensor`, *optional*):
            Y-axis heatmap predictions from the SimCC representation.
    """
    
    keypoints: torch.FloatTensor = None
    scores: torch.FloatTensor = None
    coordinate_mode: Optional[str] = None
    loss: Optional[torch.FloatTensor] = None
    pred_x: Optional[torch.FloatTensor] = None
    pred_y: Optional[torch.FloatTensor] = None


# Common layers and building blocks from RTMDet with adjustments for RTMW
class ConvModule(nn.Module):
    """A conv block that bundles conv/norm/activation layers."""
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: Union[int, Tuple[int, int]],
        stride: Union[int, Tuple[int, int]] = 1,
        padding: Union[int, Tuple[int, int]] = 0,
        dilation: Union[int, Tuple[int, int]] = 1,
        groups: int = 1,
        bias: bool = True,
        norm_cfg: Optional[Dict] = dict(type='BN'),
        act_cfg: Optional[Dict] = dict(type='SiLU'),
        inplace: bool = True,
    ):
        super().__init__()
        self.with_norm = norm_cfg is not None
        self.with_activation = act_cfg is not None
        
        # Build convolution layer
        self.conv = nn.Conv2d(
            in_channels,
            out_channels,
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=groups,
            bias=bias and not self.with_norm)
            
        # Build normalization layer
        if self.with_norm:
            norm_channels = out_channels
            # Use PyTorch default values to match MMPose's actual BN parameters during inference
            # momentum doesn't affect inference, but eps is critical!
            self.bn = nn.BatchNorm2d(norm_channels, momentum=0.1, eps=1e-05)
        
        # Build activation layer
        if self.with_activation:
            if act_cfg['type'] == 'ReLU':
                self.activate = nn.ReLU(inplace=inplace)
            elif act_cfg['type'] == 'LeakyReLU':
                self.activate = nn.LeakyReLU(negative_slope=0.1, inplace=inplace)
            elif act_cfg['type'] == 'SiLU' or act_cfg['type'] == 'Swish':
                self.activate = nn.SiLU(inplace=inplace)
            else:
                raise NotImplementedError(f"Activation {act_cfg['type']} not implemented")
            
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv(x)
        if self.with_norm:
            x = self.bn(x)
        if self.with_activation:
            x = self.activate(x)
        return x


class DepthwiseSeparableConvModule(nn.Module):
    """Depthwise separable convolution module."""
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: Union[int, Tuple[int, int]],
        stride: Union[int, Tuple[int, int]] = 1,
        padding: Union[int, Tuple[int, int]] = 0,
        dilation: Union[int, Tuple[int, int]] = 1,
        norm_cfg: Optional[Dict] = dict(type='BN'),
        act_cfg: Dict = dict(type='SiLU'),
        **kwargs
    ):
        super().__init__()
        
        # Depthwise convolution
        self.depthwise_conv = ConvModule(
            in_channels,
            in_channels,
            kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=in_channels,
            norm_cfg=norm_cfg,
            act_cfg=act_cfg,
            **kwargs)
            
        # Pointwise convolution
        self.pointwise_conv = ConvModule(
            in_channels,
            out_channels,
            1,
            norm_cfg=norm_cfg,
            act_cfg=act_cfg,
            **kwargs)
            
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.depthwise_conv(x)
        x = self.pointwise_conv(x)
        return x


class ChannelAttention(nn.Module):
    """Channel attention Module."""
    def __init__(self, channels: int) -> None:
        super().__init__()
        self.global_avgpool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Conv2d(channels, channels, 1, 1, 0, bias=True)
        self.act = nn.Hardsigmoid(inplace=True)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        with torch.amp.autocast(enabled=False, device_type=x.device.type):
            out = self.global_avgpool(x)
        out = self.fc(out)
        out = self.act(out)
        return x * out


class CSPNeXtBlock(nn.Module):
    """The basic bottleneck block used in CSPNeXt."""
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        expansion: float = 0.5,
        add_identity: bool = True,
        use_depthwise: bool = False,
        kernel_size: int = 5,
        act_cfg: Dict = dict(type='SiLU'),
    ) -> None:
        super().__init__()
        hidden_channels = int(out_channels * expansion)
        conv = DepthwiseSeparableConvModule if use_depthwise else ConvModule
        
        self.conv1 = conv(
            in_channels,
            hidden_channels,
            3,
            stride=1,
            padding=1,
            act_cfg=act_cfg)
            
        self.conv2 = DepthwiseSeparableConvModule(
            hidden_channels,
            out_channels,
            kernel_size,
            stride=1,
            padding=kernel_size // 2,
            act_cfg=act_cfg)
            
        self.add_identity = add_identity and in_channels == out_channels
            
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        identity = x
        out = self.conv1(x)
        out = self.conv2(out)
        if self.add_identity:
            return out + identity
        else:
            return out


class CSPLayer(nn.Module):
    """Cross Stage Partial Layer."""
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        expand_ratio: float = 0.5,
        num_blocks: int = 1,
        add_identity: bool = True,
        use_depthwise: bool = False,
        use_cspnext_block: bool = False,
        channel_attention: bool = False,
        act_cfg: Dict = dict(type='SiLU'),
    ) -> None:
        super().__init__()
        block = CSPNeXtBlock if use_cspnext_block else None  # Default to CSPNeXtBlock
        mid_channels = int(out_channels * expand_ratio)
        self.channel_attention = channel_attention
        
        self.main_conv = ConvModule(
            in_channels,
            mid_channels,
            1,
            act_cfg=act_cfg)
            
        self.short_conv = ConvModule(
            in_channels,
            mid_channels,
            1,
            act_cfg=act_cfg)
            
        self.final_conv = ConvModule(
            2 * mid_channels,
            out_channels,
            1,
            act_cfg=act_cfg)
            
        self.blocks = nn.Sequential(*[
            block(
                mid_channels,
                mid_channels,
                1.0,
                add_identity,
                use_depthwise) for _ in range(num_blocks)
        ])
        
        if channel_attention:
            self.attention = ChannelAttention(2 * mid_channels)
            
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x_short = self.short_conv(x)
        x_main = self.main_conv(x)
        x_main = self.blocks(x_main)
        x_final = torch.cat((x_main, x_short), dim=1)
        
        if self.channel_attention:
            x_final = self.attention(x_final)
            
        return self.final_conv(x_final)


class SPPBottleneck(nn.Module):
    """Spatial pyramid pooling layer."""
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_sizes: Tuple[int, ...] = (5, 9, 13),
        act_cfg: Dict = dict(type='SiLU'),
    ):
        super().__init__()
        mid_channels = in_channels // 2
        self.conv1 = ConvModule(
            in_channels,
            mid_channels,
            1,
            stride=1,
            act_cfg=act_cfg)
            
        self.poolings = nn.ModuleList([
            nn.MaxPool2d(kernel_size=ks, stride=1, padding=ks // 2)
            for ks in kernel_sizes
        ])
        
        conv2_channels = mid_channels * (len(kernel_sizes) + 1)
        self.conv2 = ConvModule(
            conv2_channels,
            out_channels,
            1,
            act_cfg=act_cfg)
            
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv1(x)
        with torch.amp.autocast(enabled=False, device_type=x.device.type):
            x = torch.cat(
                [x] + [pooling(x) for pooling in self.poolings], dim=1)
        x = self.conv2(x)
        return x


class CSPNeXt(nn.Module):
    """CSPNeXt backbone used in RTMW."""
    
    # From left to right:
    # in_channels, out_channels, num_blocks, add_identity, use_spp
    arch_settings = {
        'P5': [[64, 128, 3, True, False], [128, 256, 6, True, False],
               [256, 512, 6, True, False], [512, 1024, 3, False, True]],
        'P6': [[64, 128, 3, True, False], [128, 256, 6, True, False],
               [256, 512, 6, True, False], [512, 768, 3, True, False],
               [768, 1024, 3, False, True]]
    }
    
    def __init__(
        self,
        arch: str = 'P5',
        deepen_factor: float = 1.0,
        widen_factor: float = 1.0,
        out_indices: Sequence[int] = (2, 3, 4),
        frozen_stages: int = -1,
        use_depthwise: bool = False,
        expand_ratio: float = 0.5,
        channel_attention: bool = True,
        act_cfg: Dict = dict(type='SiLU'),
    ) -> None:
        super().__init__()
        arch_setting = self.arch_settings[arch]
        
        self.out_indices = out_indices
        self.frozen_stages = frozen_stages
        self.use_depthwise = use_depthwise
        
        conv = DepthwiseSeparableConvModule if use_depthwise else ConvModule
        
        self.stem = nn.Sequential(
            ConvModule(
                3,
                int(arch_setting[0][0] * widen_factor // 2),
                3,
                padding=1,
                stride=2,
                act_cfg=act_cfg),
            ConvModule(
                int(arch_setting[0][0] * widen_factor // 2),
                int(arch_setting[0][0] * widen_factor // 2),
                3,
                padding=1,
                stride=1,
                act_cfg=act_cfg),
            ConvModule(
                int(arch_setting[0][0] * widen_factor // 2),
                int(arch_setting[0][0] * widen_factor),
                3,
                padding=1,
                stride=1,
                act_cfg=act_cfg))
                
        self.layers = ['stem']
        
        for i, (in_channels, out_channels, num_blocks, add_identity,
                use_spp) in enumerate(arch_setting):
            in_channels = int(in_channels * widen_factor)
            out_channels = int(out_channels * widen_factor)
            num_blocks = max(round(num_blocks * deepen_factor), 1)
            stage = []
            
            conv_layer = conv(
                in_channels,
                out_channels,
                3,
                stride=2,
                padding=1,
                act_cfg=act_cfg)
                
            stage.append(conv_layer)
            
            if use_spp:
                spp = SPPBottleneck(
                    out_channels,
                    out_channels,
                    act_cfg=act_cfg)
                stage.append(spp)
                
            csp_layer = CSPLayer(
                out_channels,
                out_channels,
                num_blocks=num_blocks,
                add_identity=add_identity,
                use_depthwise=use_depthwise,
                use_cspnext_block=True,
                expand_ratio=expand_ratio,
                channel_attention=channel_attention,
                act_cfg=act_cfg)
                
            stage.append(csp_layer)
            
            self.add_module(f'stage{i + 1}', nn.Sequential(*stage))
            self.layers.append(f'stage{i + 1}')
    
    def freeze_stages(self) -> None:
        """Freeze stages parameters."""
        if self.frozen_stages >= 0:
            for i in range(self.frozen_stages + 1):
                m = getattr(self, self.layers[i])
                m.eval()
                for param in m.parameters():
                    param.requires_grad = False
    
    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, ...]:
        outs = []
        for i, layer_name in enumerate(self.layers):
            layer = getattr(self, layer_name)
            x = layer(x)
            if i in self.out_indices:
                outs.append(x)
        return tuple(outs)


class CSPNeXtPAFPN(nn.Module):
    """Path Aggregation Network with CSPNeXt blocks."""
    def __init__(
        self,
        in_channels: Sequence[int],
        out_channels: int,
        out_indices: Tuple[int, ...] = (1, 2),
        num_csp_blocks: int = 3,
        use_depthwise: bool = False,
        expand_ratio: float = 0.5,
        act_cfg: Dict = dict(type='SiLU'),
    ) -> None:
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.out_indices = out_indices

        conv = DepthwiseSeparableConvModule if use_depthwise else ConvModule

        # Build top-down blocks
        self.upsample = nn.Upsample(scale_factor=2, mode='nearest')
        self.reduce_layers = nn.ModuleList()
        self.top_down_blocks = nn.ModuleList()
        
        for idx in range(len(in_channels) - 1, 0, -1):
            self.reduce_layers.append(
                ConvModule(
                    in_channels[idx],
                    in_channels[idx - 1],
                    1,
                    act_cfg=act_cfg))
                    
            self.top_down_blocks.append(
                CSPLayer(
                    in_channels[idx - 1] * 2,
                    in_channels[idx - 1],
                    num_blocks=num_csp_blocks,
                    add_identity=False,
                    use_depthwise=use_depthwise,
                    use_cspnext_block=True,
                    expand_ratio=expand_ratio,
                    act_cfg=act_cfg))

        # Build bottom-up blocks
        self.downsamples = nn.ModuleList()
        self.bottom_up_blocks = nn.ModuleList()
        
        for idx in range(len(in_channels) - 1):
            self.downsamples.append(
                conv(
                    in_channels[idx],
                    in_channels[idx],
                    3,
                    stride=2,
                    padding=1,
                    act_cfg=act_cfg))
                    
            self.bottom_up_blocks.append(
                CSPLayer(
                    in_channels[idx] * 2,
                    in_channels[idx + 1],
                    num_blocks=num_csp_blocks,
                    add_identity=False,
                    use_depthwise=use_depthwise,
                    use_cspnext_block=True,
                    expand_ratio=expand_ratio,
                    act_cfg=act_cfg))

        if self.out_channels is not None:
            self.out_convs = nn.ModuleList()
            for i in range(len(in_channels)):
                self.out_convs.append(
                    conv(
                        in_channels[i],
                        out_channels,
                        3,
                        padding=1,
                        act_cfg=act_cfg))

    def forward(self, inputs: Tuple[torch.Tensor, ...]) -> Tuple[torch.Tensor, ...]:
        assert len(inputs) == len(self.in_channels)

        # Top-down path
        inner_outs = [inputs[-1]]
        for idx in range(len(self.in_channels) - 1, 0, -1):
            feat_high = inner_outs[0]
            feat_low = inputs[idx - 1]
            feat_high = self.reduce_layers[len(self.in_channels) - 1 - idx](
                feat_high)
            inner_outs[0] = feat_high

            upsample_feat = self.upsample(feat_high)

            inner_out = self.top_down_blocks[len(self.in_channels) - 1 - idx](
                torch.cat([upsample_feat, feat_low], 1))
            inner_outs.insert(0, inner_out)

        # Bottom-up path
        outs = [inner_outs[0]]
        for idx in range(len(self.in_channels) - 1):
            feat_low = outs[-1]
            feat_high = inner_outs[idx + 1]
            downsample_feat = self.downsamples[idx](feat_low)
            out = self.bottom_up_blocks[idx](
                torch.cat([downsample_feat, feat_high], 1))
            outs.append(out)

        if self.out_channels is not None:
            # Apply output convolutions
            for idx in range(len(outs)):
                outs[idx] = self.out_convs[idx](outs[idx])

        return tuple([outs[i] for i in self.out_indices])


class ScaleNorm(nn.Module):
    """Scale normalization layer with scaling factor."""
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.scale = dim ** -0.5
        self.eps = eps
        self.g = nn.Parameter(torch.ones(1))
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        norm = torch.norm(x, dim=-1, keepdim=True) * self.scale
        return x / (norm + self.eps) * self.g


class Scale(nn.Module):
    """Scale vector by element multiplications."""
    def __init__(self, dim, init_value=1., trainable=True):
        super().__init__()
        self.scale = nn.Parameter(
            init_value * torch.ones(dim), requires_grad=trainable)
        
    def forward(self, x):
        return x * self.scale


def drop_path(x: torch.Tensor,
              drop_prob: float = 0.,
              training: bool = False) -> torch.Tensor:
    """Drop paths (Stochastic Depth) per sample."""
    if drop_prob == 0. or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0], ) + (1, ) * (x.ndim - 1)
    random_tensor = keep_prob + torch.rand(
        shape, dtype=x.dtype, device=x.device)
    output = x.div(keep_prob) * random_tensor.floor()
    return output


class DropPath(nn.Module):
    """Drop paths (Stochastic Depth) per sample."""
    def __init__(self, drop_prob: float = 0.1):
        super().__init__()
        self.drop_prob = drop_prob

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return drop_path(x, self.drop_prob, self.training)


def rope(x, dim):
    """Applies Rotary Position Embedding to input tensor."""
    shape = x.shape
    if isinstance(dim, int):
        dim = [dim]
    
    spatial_shape = [shape[i] for i in dim]
    total_len = 1
    for i in spatial_shape:
        total_len *= i
        
    position = torch.reshape(
        torch.arange(total_len, dtype=torch.int, device=x.device),
        spatial_shape)
        
    for i in range(dim[-1] + 1, len(shape) - 1, 1):
        position = torch.unsqueeze(position, dim=-1)
        
    half_size = shape[-1] // 2
    freq_seq = -torch.arange(
        half_size, dtype=torch.int, device=x.device) / float(half_size)
    inv_freq = 10000**-freq_seq
    sinusoid = position[..., None] * inv_freq[None, None, :]
    sin = torch.sin(sinusoid)
    cos = torch.cos(sinusoid)
    
    x1, x2 = torch.chunk(x, 2, dim=-1)
    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)


# def gaussian_blur1d(simcc: np.ndarray, kernel: int = 11) -> np.ndarray:
#     """Modulate simcc distribution with Gaussian.

#     Note:
#         - num_keypoints: K
#         - simcc length: Wx

#     Args:
#         simcc (np.ndarray[K, Wx]): model predicted simcc.
#         kernel (int): Gaussian kernel size (K) for modulation, which should
#             match the simcc gaussian sigma when training.
#             K=17 for sigma=3 and k=11 for sigma=2.

#     Returns:
#         np.ndarray ([K, Wx]): Modulated simcc distribution.
#     """
#     assert kernel % 2 == 1

#     border = (kernel - 1) // 2
#     N, K, Wx = simcc.shape

#     for n, k in product(range(N), range(K)):
#         origin_max = np.max(simcc[n, k])
#         dr = np.zeros((1, Wx + 2 * border), dtype=np.float32)
#         dr[0, border:-border] = simcc[n, k].copy()
#         dr = cv2.GaussianBlur(dr, (kernel, 1), 0)
#         simcc[n, k] = dr[0, border:-border].copy()
#         simcc[n, k] *= origin_max / np.max(simcc[n, k])
#     return simcc

def gaussian_blur1d(simcc: torch.Tensor, kernel: int = 11) -> torch.Tensor:
    """Modulate simcc distribution with Gaussian using PyTorch.
    
    Args:
        simcc (torch.Tensor[N, K, Wx]): model predicted simcc.
        kernel (int): Gaussian kernel size (K) for modulation, which should
            match the simcc gaussian sigma when training.
            K=17 for sigma=3 and k=11 for sigma=2.

    Returns:
        torch.Tensor ([N, K, Wx]): Modulated simcc distribution.
    """
    assert kernel % 2 == 1

    border = (kernel - 1) // 2
    N, K, Wx = simcc.shape
    
    # Create Gaussian kernel
    sigma = kernel / 6.0  # Approximate conversion from kernel size to sigma
    x = torch.arange(-border, border + 1, dtype=torch.float, device=simcc.device)
    kernel_1d = torch.exp(-0.5 * (x / sigma).pow(2))
    kernel_1d = kernel_1d / kernel_1d.sum()
    
    # Reshape kernel for conv1d: (out_channels, in_channels/groups, kernel_length)
    kernel_1d = kernel_1d.view(1, 1, kernel).expand(1, 1, kernel)
    
    result = torch.zeros_like(simcc)


def get_simcc_maximum(simcc_x: torch.Tensor,
                             simcc_y: torch.Tensor,
                             apply_softmax: bool = False
                             ) -> Tuple[torch.Tensor, torch.Tensor]:
    """Get maximum response location and value from simcc representations.
    
    Args:
        simcc_x (torch.Tensor): x-axis SimCC in shape (K, Wx) or (N, K, Wx)
        simcc_y (torch.Tensor): y-axis SimCC in shape (K, Wy) or (N, K, Wy)
        apply_softmax (bool): whether to apply softmax on the heatmap.
            Defaults to False.

    Returns:
        tuple:
        - locs (torch.Tensor): locations of maximum heatmap responses in shape
            (K, 2) or (N, K, 2)
        - vals (torch.Tensor): values of maximum heatmap responses in shape
            (K,) or (N, K)
    """
    
    assert simcc_x.dim() == 2 or simcc_x.dim() == 3, f'Invalid shape {simcc_x.shape}'
    assert simcc_y.dim() == 2 or simcc_y.dim() == 3, f'Invalid shape {simcc_y.shape}'
    assert simcc_x.dim() == simcc_y.dim(), f'{simcc_x.shape} != {simcc_y.shape}'

    if simcc_x.dim() == 3:
        N, K, Wx = simcc_x.shape
        simcc_x_reshape = simcc_x.reshape(N * K, -1)
        simcc_y_reshape = simcc_y.reshape(N * K, -1)
    else:
        N = None
        simcc_x_reshape = simcc_x
        simcc_y_reshape = simcc_y

    if apply_softmax:
        simcc_x_reshape = simcc_x_reshape - torch.max(simcc_x_reshape, dim=1, keepdim=True)[0]
        simcc_y_reshape = simcc_y_reshape - torch.max(simcc_y_reshape, dim=1, keepdim=True)[0]
        ex, ey = torch.exp(simcc_x_reshape), torch.exp(simcc_y_reshape)
        simcc_x_reshape = ex / torch.sum(ex, dim=1, keepdim=True)
        simcc_y_reshape = ey / torch.sum(ey, dim=1, keepdim=True)

    # Get argmax locations
    x_locs = torch.argmax(simcc_x_reshape, dim=1)
    y_locs = torch.argmax(simcc_y_reshape, dim=1)
    
    # Create combined location tensor
    locs = torch.stack((x_locs, y_locs), dim=-1).float()
    
    # Get maximum values for each axis
    max_val_x = torch.amax(simcc_x_reshape, dim=1)
    max_val_y = torch.amax(simcc_y_reshape, dim=1)

    # Take the MINIMUM value between x and y responses (this is the correct behavior from MMPose)
    vals = torch.minimum(max_val_x, max_val_y)
    
    # Set invalid locations (where confidence is zero) to -1
    locs[vals <= 0.] = -1

    if N is not None:
        locs = locs.reshape(N, K, 2)
        vals = vals.reshape(N, K)

    return locs, vals


def refine_simcc_dark(keypoints: torch.Tensor, simcc: torch.Tensor,
                            blur_kernel_size: int) -> torch.Tensor:
    """PyTorch version of SimCC refinement using distribution aware decoding for UDP.
    
    Args:
        keypoints (torch.Tensor): The keypoint coordinates in shape (N, K, D)
        simcc (torch.Tensor): The heatmaps in shape (N, K, Wx)
        blur_kernel_size (int): The Gaussian blur kernel size of the heatmap
            modulation

    Returns:
        torch.Tensor: Refined keypoint coordinates in shape (N, K, D)
    """
    N = simcc.shape[0]

    # Modulate simcc
    simcc = gaussian_blur1d(simcc, blur_kernel_size)
    simcc = torch.clamp(simcc, min=1e-3, max=50.)
    simcc = torch.log(simcc)

    # Pad the simcc tensor
    simcc = F.pad(simcc, (2, 2), mode='replicate')

    # Create refined keypoints tensor
    keypoints_refined = keypoints.clone()

    for n in range(N):
        # Convert keypoints to indices
        px = (keypoints[n] + 2.5).long().view(-1, 1)  # K, 1
        
        # Ensure indices are within bounds
        px = torch.clamp(px, min=0, max=simcc.shape[2]-1)
        
        # Sample values for dx calculation
        # Use gather for more efficient tensor indexing
        # Create index tensors for gather
        batch_idx = torch.zeros_like(px).long() + n
        channel_idx = torch.arange(px.shape[0], device=px.device).view(-1, 1)
        
        # Gather values for dx and dxx calculation
        dx0 = simcc[n, torch.arange(px.shape[0], device=px.device), px.squeeze(-1)]
        dx1 = simcc[n, torch.arange(px.shape[0], device=px.device), (px + 1).squeeze(-1)]
        dx_1 = simcc[n, torch.arange(px.shape[0], device=px.device), (px - 1).squeeze(-1)]
        dx2 = simcc[n, torch.arange(px.shape[0], device=px.device), (px + 2).squeeze(-1)]
        dx_2 = simcc[n, torch.arange(px.shape[0], device=px.device), (px - 2).squeeze(-1)]
        
        # Calculate dx and dxx
        dx = 0.5 * (dx1 - dx_1)
        dxx = 1e-9 + 0.25 * (dx2 - 2 * dx0 + dx_2)
        
        # Calculate offset
        offset = dx / dxx
        
        # Apply offset to refine keypoints
        keypoints_refined[n] -= offset

    return keypoints_refined


class SimCCCodec:
    """Generate keypoint representation via SimCC approach - All PyTorch implementation.
    
    This class implements the SimCC (Simple Coordinate Classification) approach for human pose estimation
    without relying on NumPy, ensuring full PyTorch tensor compatibility.
    
    Args:
        input_size (tuple): Input image size in [w, h]
        smoothing_type (str): The SimCC label smoothing strategy. Options are
            'gaussian' and 'standard'. Defaults to 'gaussian'
        sigma (float | int | tuple): The sigma value in the Gaussian SimCC label.
            Defaults to 6.0
        simcc_split_ratio (float): The ratio of the label size to the input size.
            For example, if the input width is w, the x label size will be
            w*simcc_split_ratio. Defaults to 2.0
        normalize (bool): Whether to normalize the heatmaps. Defaults to False.
        use_dark (bool): Whether to use the DARK post processing. Defaults to False.
    """
    
    def __init__(
        self,
        input_size,
        smoothing_type='gaussian',
        sigma=6.0,
        simcc_split_ratio=2.0,
        normalize=False,
        use_dark=False
    ):
        self.input_size = input_size
        self.smoothing_type = smoothing_type
        self.simcc_split_ratio = simcc_split_ratio
        self.normalize = normalize
        self.use_dark = use_dark
        
        if isinstance(sigma, (float, int)):
            sigma = [sigma, sigma]
        self.sigma = torch.tensor(sigma)
    
    def encode(self, keypoints, keypoints_visible=None):
        """Encoding keypoints into SimCC labels. Note that the original
        keypoint coordinates should be in the input image space.
        
        This is primarily used for training but included for completeness.
        """
        raise NotImplementedError(
            "SimCCCodecPyTorch.encode() is not implemented, only supports inference.")
    
    def decode(self, simcc_x: torch.Tensor,
              simcc_y: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Decode keypoint coordinates from SimCC representations. The decoded
        coordinates are in the input image space.

        Args:
            simcc_x (torch.Tensor): SimCC label for x-axis
            simcc_y (torch.Tensor): SimCC label for y-axis

        Returns:
            tuple:
            - keypoints (torch.Tensor): Decoded coordinates in shape (N, K, D)
            - scores (torch.Tensor): The keypoint scores in shape (N, K).
                It usually represents the confidence of the keypoint prediction
        """
        device = simcc_x.device
        
        # Ensure correct dimensions for processing
        if simcc_x.dim() == 2:
            simcc_x = simcc_x.unsqueeze(0)  # Add batch dimension
        if simcc_y.dim() == 2:
            simcc_y = simcc_y.unsqueeze(0)  # Add batch dimension
            
        keypoints, scores = get_simcc_maximum(simcc_x, simcc_y)

        # Apply DARK post-processing if requested
        if self.use_dark:
            # Calculate blur kernel sizes based on sigma values
            sigma_tensor = self.sigma.to(device)
            x_blur = int((sigma_tensor[0] * 20 - 7) // 3)
            y_blur = int((sigma_tensor[1] * 20 - 7) // 3)
            
            # Ensure odd kernel sizes
            x_blur -= int((x_blur % 2) == 0)
            y_blur -= int((y_blur % 2) == 0)
            
            # Apply DARK refinement separately to x and y coordinates
            for i in range(keypoints.shape[0]):
                keypoints_x = keypoints[i, :, 0:1]
                keypoints_y = keypoints[i, :, 1:2]
                
                keypoints[i, :, 0] = refine_simcc_dark(
                    keypoints_x, simcc_x[i:i+1], x_blur)[:, 0]
                keypoints[i, :, 1] = refine_simcc_dark(
                    keypoints_y, simcc_y[i:i+1], y_blur)[:, 0]

        # Convert from SimCC coordinate space back to image coordinate space
        keypoints /= self.simcc_split_ratio

        return keypoints, scores


class RTMCCBlock(nn.Module):
    """Gated Attention Unit (GAU) in RTMBlock."""
    
    def __init__(
        self,
        num_token,
        in_token_dims,
        out_token_dims,
        expansion_factor=2,
        s=128,
        eps=1e-5,
        dropout_rate=0.,
        drop_path=0.,
        attn_type='self-attn',
        act_fn='SiLU',
        bias=False,
        use_rel_bias=True,
        pos_enc=False
    ):
        super(RTMCCBlock, self).__init__()
        self.s = s
        self.num_token = num_token
        self.use_rel_bias = use_rel_bias
        self.attn_type = attn_type
        self.pos_enc = pos_enc
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.e = int(in_token_dims * expansion_factor)
        
        if use_rel_bias:
            if attn_type == 'self-attn':
                self.w = nn.Parameter(
                    torch.rand([2 * num_token - 1], dtype=torch.float))
            else:
                self.a = nn.Parameter(torch.rand([1, s], dtype=torch.float))
                self.b = nn.Parameter(torch.rand([1, s], dtype=torch.float))
                
        self.o = nn.Linear(self.e, out_token_dims, bias=bias)
        
        if attn_type == 'self-attn':
            self.uv = nn.Linear(in_token_dims, 2 * self.e + self.s, bias=bias)
            self.gamma = nn.Parameter(torch.rand((2, self.s)))
            self.beta = nn.Parameter(torch.rand((2, self.s)))
        else:
            self.uv = nn.Linear(in_token_dims, self.e + self.s, bias=bias)
            self.k_fc = nn.Linear(in_token_dims, self.s, bias=bias)
            self.v_fc = nn.Linear(in_token_dims, self.e, bias=bias)
            nn.init.xavier_uniform_(self.k_fc.weight)
            nn.init.xavier_uniform_(self.v_fc.weight)
            
        self.ln = ScaleNorm(in_token_dims, eps=eps)
        nn.init.xavier_uniform_(self.uv.weight)
        
        if act_fn == 'SiLU' or act_fn == nn.SiLU:
            self.act_fn = nn.SiLU(True)
        elif act_fn == 'ReLU' or act_fn == nn.ReLU:
            self.act_fn = nn.ReLU(True)
        else:
            raise NotImplementedError
            
        if in_token_dims == out_token_dims:
            self.shortcut = True
            self.res_scale = Scale(in_token_dims)
        else:
            self.shortcut = False
            
        self.register_buffer('sqrt_s', torch.sqrt(torch.tensor(s, dtype=torch.float)), persistent=False)
        self.dropout_rate = dropout_rate
        if dropout_rate > 0.:
            self.dropout = nn.Dropout(dropout_rate)
            
    def rel_pos_bias(self, seq_len, k_len=None):
        """Add relative position bias."""
        if self.attn_type == 'self-attn':
            t = F.pad(self.w[:2 * seq_len - 1], [0, seq_len]).repeat(seq_len)
            t = t[..., :-seq_len].reshape(-1, seq_len, 3 * seq_len - 2)
            r = (2 * seq_len - 1) // 2
            t = t[..., r:-r]
        else:
            a = rope(self.a.repeat(seq_len, 1), dim=0)
            b = rope(self.b.repeat(k_len, 1), dim=0)
            t = torch.bmm(a, b.permute(0, 2, 1))
        return t
        
    def _forward(self, inputs):
        """GAU Forward function."""
        if self.attn_type == 'self-attn':
            x = inputs
        else:
            x, k, v = inputs
            
        x = self.ln(x)
        uv = self.uv(x)
        uv = self.act_fn(uv)
        
        if self.attn_type == 'self-attn':
            # Split into u, v, base
            u, v, base = torch.split(uv, [self.e, self.e, self.s], dim=2)
            # Apply gamma and beta parameters
            base = base.unsqueeze(2) * self.gamma[None, None, :] + self.beta[None, None, :]
            if self.pos_enc:
                base = rope(base, dim=1)
            # Split base into q, k
            q, k = torch.unbind(base, dim=2)
        else:
            # Split into u, q
            u, q = torch.split(uv, [self.e, self.s], dim=2)
            k = self.k_fc(k)  # -> [B, K, s]
            v = self.v_fc(v)  # -> [B, K, e]
            if self.pos_enc:
                q = rope(q, 1)
                k = rope(k, 1)
                
        # Calculate attention
        qk = torch.bmm(q, k.permute(0, 2, 1))
        
        if self.use_rel_bias:
            if self.attn_type == 'self-attn':
                bias = self.rel_pos_bias(q.size(1))
            else:
                bias = self.rel_pos_bias(q.size(1), k.size(1))
            qk += bias[:, :q.size(1), :k.size(1)]
            
        # Apply kernel (square of ReLU)
        kernel = torch.square(F.relu(qk / self.sqrt_s))
        
        if self.dropout_rate > 0.:
            kernel = self.dropout(kernel)
            
        # Apply attention
        if self.attn_type == 'self-attn':
            x = u * torch.bmm(kernel, v)
        else:
            x = u * torch.bmm(kernel, v)
            
        x = self.o(x)
        return x
        
    def forward(self, x):
        """Forward function."""
        if self.shortcut:
            if self.attn_type == 'cross-attn':
                res_shortcut = x[0]
            else:
                res_shortcut = x
            main_branch = self.drop_path(self._forward(x))
            return self.res_scale(res_shortcut) + main_branch
        else:
            return self.drop_path(self._forward(x))


class RTMWHead(nn.Module):
    """Top-down head introduced in RTMPose-Wholebody (2023).
    Updated to use PyTorch-only implementations without NumPy or OpenCV.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        input_size: Tuple[int, int],
        in_featuremap_size: Tuple[int, int],
        simcc_split_ratio: float = 2.0,
        final_layer_kernel_size: int = 7,
        gau_cfg: Optional[Dict] = None,
        decoder: Optional[Dict] = None,
    ):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.input_size = input_size
        self.in_featuremap_size = in_featuremap_size
        self.simcc_split_ratio = simcc_split_ratio

        # Default GAU config if not provided
        if gau_cfg is None:
            gau_cfg = dict(
                hidden_dims=256,
                s=128,
                expansion_factor=2,
                dropout_rate=0.,
                drop_path=0.,
                act_fn='ReLU',
                use_rel_bias=False,
                pos_enc=False)

        # Define SimCC layers
        flatten_dims = self.in_featuremap_size[0] * self.in_featuremap_size[1]

        ps = 2  # pixel shuffle factor
        self.ps = nn.PixelShuffle(ps)
        
        self.conv_dec = ConvModule(
            in_channels // ps**2,
            in_channels // 4,
            kernel_size=final_layer_kernel_size,
            stride=1,
            padding=final_layer_kernel_size // 2,
            norm_cfg=dict(type='BN'),
            act_cfg=dict(type='ReLU'))

        self.final_layer = ConvModule(
            in_channels,
            out_channels,
            kernel_size=final_layer_kernel_size,
            stride=1,
            padding=final_layer_kernel_size // 2,
            norm_cfg=dict(type='BN'),
            act_cfg=dict(type='ReLU'))
            
        self.final_layer2 = ConvModule(
            in_channels // ps + in_channels // 4,
            out_channels,
            kernel_size=final_layer_kernel_size,
            stride=1,
            padding=final_layer_kernel_size // 2,
            norm_cfg=dict(type='BN'),
            act_cfg=dict(type='ReLU'))

        self.mlp = nn.Sequential(
            ScaleNorm(flatten_dims),
            nn.Linear(flatten_dims, gau_cfg['hidden_dims'] // 2, bias=False))

        self.mlp2 = nn.Sequential(
            ScaleNorm(flatten_dims * ps**2),
            nn.Linear(
                flatten_dims * ps**2, gau_cfg['hidden_dims'] // 2, bias=False))

        W = int(self.input_size[0] * self.simcc_split_ratio)
        H = int(self.input_size[1] * self.simcc_split_ratio)

        self.gau = RTMCCBlock(
            self.out_channels,
            gau_cfg['hidden_dims'],
            gau_cfg['hidden_dims'],
            s=gau_cfg['s'],
            expansion_factor=gau_cfg['expansion_factor'],
            dropout_rate=gau_cfg['dropout_rate'],
            drop_path=gau_cfg['drop_path'],
            attn_type='self-attn',
            act_fn=gau_cfg['act_fn'],
            use_rel_bias=gau_cfg['use_rel_bias'],
            pos_enc=gau_cfg['pos_enc'])

        self.cls_x = nn.Linear(gau_cfg['hidden_dims'], W, bias=False)
        self.cls_y = nn.Linear(gau_cfg['hidden_dims'], H, bias=False)

        # Create SimCC codec for decoding - using PyTorch version
        if decoder is not None:
            self.decoder = SimCCCodec(
                input_size=decoder.get('input_size', self.input_size),
                smoothing_type=decoder.get('smoothing_type', 'gaussian'),
                sigma=decoder.get('sigma', (4.9, 5.66)),
                simcc_split_ratio=self.simcc_split_ratio,
                normalize=decoder.get('normalize', False),
                use_dark=decoder.get('use_dark', False)
            )
        else:
            self.decoder = SimCCCodec(
                input_size=self.input_size,
                sigma=(4.9, 5.66),
                simcc_split_ratio=self.simcc_split_ratio,
                normalize=False,
                use_dark=False
            )

    def forward(self, feats: Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
        """Forward the network to get SimCC representations.

        Args:
            feats (Tuple[Tensor]): Multi scale feature maps.

        Returns:
            pred_x (Tensor): 1d representation of x.
            pred_y (Tensor): 1d representation of y.
        """
        enc_b, enc_t = feats

        feats_t = self.final_layer(enc_t)
        feats_t = torch.flatten(feats_t, 2)
        feats_t = self.mlp(feats_t)

        dec_t = self.ps(enc_t)
        dec_t = self.conv_dec(dec_t)
        enc_b = torch.cat([dec_t, enc_b], dim=1)

        feats_b = self.final_layer2(enc_b)
        feats_b = torch.flatten(feats_b, 2)
        feats_b = self.mlp2(feats_b)

        feats = torch.cat([feats_t, feats_b], dim=2)

        feats = self.gau(feats)

        pred_x = self.cls_x(feats)
        pred_y = self.cls_y(feats)

        return pred_x, pred_y
        
    def predict(self, feats: Tuple[torch.Tensor, torch.Tensor], flip_test=False, flip_indices=None):
        """Predict keypoints from features.
        
        Args:
            feats (Tuple[torch.Tensor]): Features from the backbone + neck
            flip_test (bool): Whether to use flip test augmentation
            flip_indices (List[int]): Indices for flipping keypoints
            
        Returns:
            List[Dict]: Predicted keypoints and scores
        """
        batch_pred_x, batch_pred_y = None, None
        device = feats[0].device
        
        if flip_test:
            assert flip_indices is not None, "flip_indices must be provided for flip test"
            
            # Original forward pass
            _batch_pred_x, _batch_pred_y = self.forward(feats)
            
            # Create flipped input and get predictions
            feats_flipped = [torch.flip(feat, dims=[-1]) for feat in feats]
            _batch_pred_x_flip, _batch_pred_y_flip = self.forward(feats_flipped)
            
            # Flip predictions back - critical part
            _batch_pred_x_flip = torch.flip(_batch_pred_x_flip, dims=[2])  # Flip along the width dimension
            
            # Handle keypoint swapping (like left-right joints)
            batch_size = _batch_pred_x.shape[0]
            for i in range(batch_size):
                for src_idx, dst_idx in enumerate(flip_indices):
                    if src_idx != dst_idx:
                        _batch_pred_x_flip[i, dst_idx] = _batch_pred_x_flip[i, src_idx].clone()
                        _batch_pred_y_flip[i, dst_idx] = _batch_pred_y_flip[i, src_idx].clone()
            
            # Average the predictions
            batch_pred_x = (_batch_pred_x + _batch_pred_x_flip) * 0.5
            batch_pred_y = (_batch_pred_y + _batch_pred_y_flip) * 0.5
        else:
            # Standard forward pass
            batch_pred_x, batch_pred_y = self.forward(feats)
        
        # Decode keypoints using PyTorch-based decoder
        keypoints, scores = self.decoder.decode(batch_pred_x, batch_pred_y)
        
        # Convert to list of instances
        batch_size = keypoints.shape[0]
        instances = []
        
        for i in range(batch_size):
            instances.append({
                'keypoints': keypoints[i],
                'keypoint_scores': scores[i]
            })
            
        return instances


class RTMWModel(PreTrainedModel):
    """
    RTMW model for human pose estimation.
    
    This model consists of a backbone, neck, and pose head for keypoint detection.
    All implementations use PyTorch only with no NumPy or OpenCV dependencies.
    """

    config_class = RTMWConfig
    base_model_prefix = "rtmw"
    main_input_name = "pixel_values"
    
    def __init__(self, config: RTMWConfig):
        super().__init__(config)
        self.config = config
        
        self.backbone = CSPNeXt(
            arch=config.backbone_arch,
            deepen_factor=config.backbone_deepen_factor,
            widen_factor=config.backbone_widen_factor,
            expand_ratio=config.backbone_expand_ratio,
            channel_attention=config.backbone_channel_attention,
            use_depthwise=False,
        )
        
        # Build neck
        self.neck = CSPNeXtPAFPN(
            in_channels=config.neck_in_channels,
            out_channels=config.neck_out_channels,
            num_csp_blocks=config.neck_num_csp_blocks,
            expand_ratio=config.neck_expand_ratio,
            use_depthwise=False,
        )
        
        # Build head
        # Create GAU config from the configuration
        gau_cfg = {
            'hidden_dims': config.gau_hidden_dims,
            's': config.gau_s,
            'expansion_factor': config.gau_expansion_factor,
            'dropout_rate': config.gau_dropout_rate,
            'drop_path': config.gau_drop_path,
            'act_fn': config.gau_act_fn,
            'use_rel_bias': config.gau_use_rel_bias,
            'pos_enc': config.gau_pos_enc,
        }
        
        self.head = RTMWHead(
            in_channels=config.head_in_channels,
            out_channels=config.num_keypoints,
            input_size=config.input_size,
            in_featuremap_size=config.head_in_featuremap_size,
            simcc_split_ratio=config.simcc_split_ratio,
            final_layer_kernel_size=config.head_final_layer_kernel_size,
            gau_cfg=gau_cfg,
            decoder = dict(
                input_size=config.input_size,
                sigma=config.decoder_sigma,
                simcc_split_ratio=config.simcc_split_ratio,
                normalize=config.decoder_normalize,
                use_dark=config.decoder_use_dark)
        )
        
        # Initialize weights
        self.init_weights()
        # Required: triggers post_init() which sets all_tied_weights_keys etc.
        self.post_init()
    
    def init_weights(self):
        """Initialize the weights of the model."""
        # Initialize convolution layers with normal distribution
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.normal_(m.weight, mean=0, std=0.01)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            if isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            if isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, mean=0, std=0.01)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
    
    def forward(
        self,
        pixel_values=None,
        bbox=None,
        coordinate_mode: str = "image",
        labels=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        """
        Forward pass of the model.
        
        Args:
            pixel_values (`torch.FloatTensor` of shape `(batch_size, channels, height, width)`):
                Pixel values cropped and resized to the model's input resolution
                (e.g. 288×384).  Use `RTMWImageProcessor` or prepare manually with
                ImageNet normalisation.
            bbox (`torch.FloatTensor` of shape `(batch_size, 4)` or `(4,)`, *optional*):
                Person bounding boxes in the **original** image, as
                ``[x1, y1, x2, y2]`` pixel coordinates.  Required when
                ``coordinate_mode="image"``; ignored otherwise.
            coordinate_mode (`str`, *optional*, defaults to ``"image"``):
                How to express the returned keypoint coordinates:

                - ``"model"``         — raw SimCC space (same resolution as the
                  model input, e.g. 288×384 px).  No extra arguments needed.
                - ``"image"``         — rescaled back to the original image pixel
                  space using the supplied ``bbox``.  If ``bbox`` is ``None`` the
                  output falls back to ``"model"`` space with a warning.
                - ``"root_relative"`` — root-normalised coordinates.  The root is
                  the midpoint of the left-hip (kp 11) and right-hip (kp 12)
                  joints.  All keypoints are translated so the root is at the
                  origin, then divided by half the inter-hip distance so that
                  each hip lands at unit distance from the origin.  Applied
                  *after* any ``"image"`` projection when both are combined
                  (not combinable via this single arg — choose one).
            labels (`List[Dict]`, *optional*):
                Labels for computing the pose estimation loss.
            output_hidden_states (`bool`, *optional*):
                Whether or not to return the hidden states of all layers.
            return_dict (`bool`, *optional*):
                Whether or not to return a ModelOutput instead of a plain tuple.
                
        Returns:
            `PoseOutput` or `tuple`:
                If return_dict=True, `PoseOutput` is returned.
                If return_dict=False, a tuple is returned with keypoints and scores.
        """
        import warnings
        return_dict = return_dict if return_dict is not None else True
        
        # Get inputs
        if pixel_values is None:
            raise ValueError("You have to specify pixel_values")
            
        # Extract features from backbone
        backbone_features = self.backbone(pixel_values)
        
        # Process features through neck
        neck_features = self.neck(backbone_features)
        
        # Get SimCC representations from pose head
        pred_x, pred_y = self.head.forward(neck_features)
        
        # Decode keypoints
        instances = self.head.predict(neck_features, None)
        
        # Extract keypoints and scores from instances
        batch_size = len(instances)
        keypoints = torch.zeros((batch_size, self.head.out_channels, 2), device=pixel_values.device)
        scores = torch.zeros((batch_size, self.head.out_channels), device=pixel_values.device)
        
        for i, instance in enumerate(instances):
            keypoints[i] = instance['keypoints']
            scores[i] = instance['keypoint_scores']

        # Apply fixed min-max normalization to map scores to [0, 1].
        # Only valid scores (> 0) are normalized; invalid keypoints keep
        # their raw (≤ 0) values so downstream code can still filter them.
        score_min = getattr(self.config, 'score_min', None)
        score_max = getattr(self.config, 'score_max', None)
        if score_min is not None and score_max is not None and score_max > score_min:
            valid_mask = scores > 0
            scores[valid_mask] = torch.clamp(
                (scores[valid_mask] - score_min) / (score_max - score_min),
                0.0, 1.0,
            )

        # ── Coordinate transform ──────────────────────────────────────────────
        # Keypoints are currently in model-input space:
        #   x in [0, model_w),  y in [0, model_h)
        # e.g. model_w=288, model_h=384 for rtmw-l-384x288.
        if coordinate_mode == "image":
            if bbox is None:
                warnings.warn(
                    "coordinate_mode='image' requires bbox=[x1,y1,x2,y2] per image. "
                    "Falling back to model-space coordinates.",
                    UserWarning, stacklevel=2,
                )
                coordinate_mode = "model"
            else:
                # bbox: (B, 4) or (4,) → normalise to (B, 1, 2) broadcast shape
                bbox_t = torch.as_tensor(bbox, dtype=keypoints.dtype, device=keypoints.device)
                if bbox_t.dim() == 1:
                    bbox_t = bbox_t.unsqueeze(0).expand(keypoints.shape[0], -1)
                model_h = pixel_values.shape[2]   # H dim of model input
                model_w = pixel_values.shape[3]   # W dim of model input
                x1 = bbox_t[:, 0:1]   # (B, 1)
                y1 = bbox_t[:, 1:2]
                x2 = bbox_t[:, 2:3]
                y2 = bbox_t[:, 3:4]
                scale_x = (x2 - x1) / model_w    # (B, 1)
                scale_y = (y2 - y1) / model_h    # (B, 1)
                # (B, K, 2) — broadcast over K
                keypoints = keypoints.clone()
                keypoints[:, :, 0] = keypoints[:, :, 0] * scale_x + x1
                keypoints[:, :, 1] = keypoints[:, :, 1] * scale_y + y1

        elif coordinate_mode == "root_relative":
            # Root = midpoint of left_hip (11) and right_hip (12).
            # Scale = half the inter-hip distance so each hip is at unit
            # distance from the root.  Clamp to ≥1 px to guard against
            # degenerate detections where the hips are co-located.
            left_hip  = keypoints[:, 11, :]                             # (B, 2)
            right_hip = keypoints[:, 12, :]                             # (B, 2)
            root      = 0.5 * (left_hip + right_hip)                    # (B, 2)
            scale     = (0.5 * torch.norm(right_hip - left_hip, dim=-1, keepdim=True)  # (B, 1)
                         .clamp(min=1.0))
            keypoints = (keypoints - root.unsqueeze(1)) / scale.unsqueeze(1)

        elif coordinate_mode != "model":
            raise ValueError(
                f"coordinate_mode must be 'model', 'image', or 'root_relative', got {coordinate_mode!r}"
            )
        # ─────────────────────────────────────────────────────────────────────

        if return_dict:
            return PoseOutput(
                keypoints=keypoints,
                scores=scores,
                coordinate_mode=coordinate_mode,
                pred_x=pred_x,
                pred_y=pred_y
            )
        else:
            return (keypoints, scores)