"""
Copied from RT-DETR (https://github.com/lyuwenyu/RT-DETR)
Copyright(c) 2023 lyuwenyu. All Rights Reserved.
"""

import torch
import torch.nn as nn

from ...core import register

__all__ = ["Classification", "ClassHead"]


@register()
class Classification(torch.nn.Module):
    __inject__ = ["backbone", "head"]

    def __init__(self, backbone: nn.Module, head: nn.Module = None):
        super().__init__()

        self.backbone = backbone
        self.head = head

    def forward(self, x):
        x = self.backbone(x)

        if self.head is not None:
            x = self.head(x)

        return x


@register()
class ClassHead(nn.Module):
    def __init__(self, hidden_dim, num_classes):
        super().__init__()
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.proj = nn.Linear(hidden_dim, num_classes)

    def forward(self, x):
        x = x[0] if isinstance(x, (list, tuple)) else x
        x = self.pool(x)
        x = x.reshape(x.shape[0], -1)
        x = self.proj(x)
        return x