head_extractor/src/mmengine/utils/dl_utils/misc.py

# Copyright (c) OpenMMLab. All rights reserved.
import pkgutil
from typing import Optional, Tuple, Union

import numpy as np
import torch
import torch.nn as nn

from ..misc import is_tuple_of
from .parrots_wrapper import _BatchNorm, _InstanceNorm


def is_norm(layer: nn.Module,
            exclude: Optional[Union[type, Tuple[type]]] = None) -> bool:
    """Check if a layer is a normalization layer.

    Args:
        layer (nn.Module): The layer to be checked.
        exclude (type, tuple[type], optional): Types to be excluded.

    Returns:
        bool: Whether the layer is a norm layer.
    """
    if exclude is not None:
        if not isinstance(exclude, tuple):
            exclude = (exclude, )
        if not is_tuple_of(exclude, type):
            raise TypeError(
                f'"exclude" must be either None or type or a tuple of types, '
                f'but got {type(exclude)}: {exclude}')

    if exclude and isinstance(layer, exclude):
        return False

    all_norm_bases = (_BatchNorm, _InstanceNorm, nn.GroupNorm, nn.LayerNorm)
    return isinstance(layer, all_norm_bases)


def tensor2imgs(tensor: torch.Tensor,
                mean: Optional[Tuple[float, float, float]] = None,
                std: Optional[Tuple[float, float, float]] = None,
                to_bgr: bool = True):
    """Convert tensor to 3-channel images or 1-channel gray images.

    Args:
        tensor (torch.Tensor): Tensor that contains multiple images, shape (
            N, C, H, W). :math:`C` can be either 3 or 1. If C is 3, the format
            should be RGB.
        mean (tuple[float], optional): Mean of images. If None,
            (0, 0, 0) will be used for tensor with 3-channel,
            while (0, ) for tensor with 1-channel. Defaults to None.
        std (tuple[float], optional): Standard deviation of images. If None,
            (1, 1, 1) will be used for tensor with 3-channel,
            while (1, ) for tensor with 1-channel. Defaults to None.
        to_bgr (bool): For the tensor with 3 channel, convert its format to
            BGR. For the tensor with 1 channel, it must be False. Defaults to
            True.

    Returns:
        list[np.ndarray]: A list that contains multiple images.
    """

    assert torch.is_tensor(tensor) and tensor.ndim == 4
    channels = tensor.size(1)
    assert channels in [1, 3]
    if mean is None:
        mean = (0, ) * channels
    if std is None:
        std = (1, ) * channels
    assert (channels == len(mean) == len(std) == 3) or \
           (channels == len(mean) == len(std) == 1 and not to_bgr)
    mean = tensor.new_tensor(mean).view(1, -1)
    std = tensor.new_tensor(std).view(1, -1)
    tensor = tensor.permute(0, 2, 3, 1) * std + mean
    imgs = tensor.detach().cpu().numpy()
    if to_bgr and channels == 3:
        imgs = imgs[:, :, :, (2, 1, 0)]  # RGB2BGR
    imgs = [np.ascontiguousarray(img) for img in imgs]
    return imgs


def has_batch_norm(model: nn.Module) -> bool:
    """Detect whether model has a BatchNormalization layer.

    Args:
        model (nn.Module): training model.

    Returns:
        bool: whether model has a BatchNormalization layer
    """
    if isinstance(model, _BatchNorm):
        return True
    for m in model.children():
        if has_batch_norm(m):
            return True
    return False


def mmcv_full_available() -> bool:
    """Check whether mmcv-full is installed.

    Returns:
        bool: True if mmcv-full is installed else False.
    """
    try:
        import mmcv  # noqa: F401
    except ImportError:
        return False
    ext_loader = pkgutil.find_loader('mmcv._ext')
    return ext_loader is not None