"""上采样 neck：把骨架的低分辨率特征恢复到定位所需的 stride。"""

from __future__ import annotations

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

from models.layers import ConvBNAct
from models.registry import NECKS


@NECKS.register("upsample")
class UpsampleNeck(nn.Module):
    """`num_upsamples` 次 2 倍上采样，每次接一个 3x3 卷积平滑。

    使用 nearest + conv 而非转置卷积：避免棋盘伪影，且在推理引擎上更快。
    """

    def __init__(self, in_channels: int, channels: int = 64, num_upsamples: int = 2) -> None:
        super().__init__()
        self.lateral = ConvBNAct(in_channels, channels, kernel_size=1)
        self.blocks = nn.ModuleList(
            [ConvBNAct(channels, channels, kernel_size=3) for _ in range(num_upsamples)]
        )
        self.out_channels = channels
        self.scale_factor = 2**num_upsamples

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.lateral(x)
        for block in self.blocks:
            x = block(F.interpolate(x, scale_factor=2.0, mode="nearest"))
        return x