# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
from mmengine.model import BaseModule

from mmdet.registry import MODELS


@MODELS.register_module()
class GlobalAveragePooling(BaseModule):
    """Global Average Pooling neck.

    Note that we use `view` to remove extra channel after pooling. We do not
    use `squeeze` as it will also remove the batch dimension when the tensor
    has a batch dimension of size 1, which can lead to unexpected errors.
    """

    def __init__(self, kernel_size=None, stride=None):
        super(GlobalAveragePooling, self).__init__()
        if kernel_size is None and stride is None:
            self.gap = nn.AdaptiveAvgPool2d((1, 1))
        else:
            self.gap = nn.AvgPool2d(kernel_size, stride)

    def forward(self, inputs):
        if isinstance(inputs, tuple):
            outs = tuple([self.gap(x) for x in inputs])
            outs = tuple([
                out.view(x.size(0),
                         torch.tensor(out.size()[1:]).prod())
                for out, x in zip(outs, inputs)
            ])
        elif isinstance(inputs, torch.Tensor):
            outs = self.gap(inputs)
            outs = outs.view(
                inputs.size(0),
                torch.tensor(outs.size()[1:]).prod())
        else:
            raise TypeError('neck inputs should be tuple or torch.tensor')
        return outs