mmyolo/utils/boxam_utils.py

# Copyright (c) OpenMMLab. All rights reserved.
import bisect
import copy
import warnings
from pathlib import Path
from typing import Callable, List, Optional, Tuple, Union

import cv2
import numpy as np
import torch
import torch.nn as nn
import torchvision
from mmcv.transforms import Compose
from mmdet.evaluation import get_classes
from mmdet.utils import ConfigType
from mmengine.config import Config
from mmengine.registry import init_default_scope
from mmengine.runner import load_checkpoint
from mmengine.structures import InstanceData
from torch import Tensor

from mmyolo.registry import MODELS

try:
    from pytorch_grad_cam import (AblationCAM, AblationLayer,
                                  ActivationsAndGradients)
    from pytorch_grad_cam import GradCAM as Base_GradCAM
    from pytorch_grad_cam import GradCAMPlusPlus as Base_GradCAMPlusPlus
    from pytorch_grad_cam.base_cam import BaseCAM
    from pytorch_grad_cam.utils.image import scale_cam_image, show_cam_on_image
    from pytorch_grad_cam.utils.svd_on_activations import get_2d_projection
except ImportError:
    pass


def init_detector(
    config: Union[str, Path, Config],
    checkpoint: Optional[str] = None,
    palette: str = 'coco',
    device: str = 'cuda:0',
    cfg_options: Optional[dict] = None,
) -> nn.Module:
    """Initialize a detector from config file.

    Args:
        config (str, :obj:`Path`, or :obj:`mmengine.Config`): Config file path,
            :obj:`Path`, or the config object.
        checkpoint (str, optional): Checkpoint path. If left as None, the model
            will not load any weights.
        palette (str): Color palette used for visualization. If palette
            is stored in checkpoint, use checkpoint's palette first, otherwise
            use externally passed palette. Currently, supports 'coco', 'voc',
            'citys' and 'random'. Defaults to coco.
        device (str): The device where the anchors will be put on.
            Defaults to cuda:0.
        cfg_options (dict, optional): Options to override some settings in
            the used config.

    Returns:
        nn.Module: The constructed detector.
    """
    if isinstance(config, (str, Path)):
        config = Config.fromfile(config)
    elif not isinstance(config, Config):
        raise TypeError('config must be a filename or Config object, '
                        f'but got {type(config)}')
    if cfg_options is not None:
        config.merge_from_dict(cfg_options)
    elif 'init_cfg' in config.model.backbone:
        config.model.backbone.init_cfg = None

    # only change this
    # grad based method requires train_cfg
    # config.model.train_cfg = None
    init_default_scope(config.get('default_scope', 'mmyolo'))

    model = MODELS.build(config.model)
    if checkpoint is not None:
        checkpoint = load_checkpoint(model, checkpoint, map_location='cpu')
        # Weights converted from elsewhere may not have meta fields.
        checkpoint_meta = checkpoint.get('meta', {})
        # save the dataset_meta in the model for convenience
        if 'dataset_meta' in checkpoint_meta:
            # mmdet 3.x, all keys should be lowercase
            model.dataset_meta = {
                k.lower(): v
                for k, v in checkpoint_meta['dataset_meta'].items()
            }
        elif 'CLASSES' in checkpoint_meta:
            # < mmdet 3.x
            classes = checkpoint_meta['CLASSES']
            model.dataset_meta = {'classes': classes, 'palette': palette}
        else:
            warnings.simplefilter('once')
            warnings.warn(
                'dataset_meta or class names are not saved in the '
                'checkpoint\'s meta data, use COCO classes by default.')
            model.dataset_meta = {
                'classes': get_classes('coco'),
                'palette': palette
            }

    model.cfg = config  # save the config in the model for convenience
    model.to(device)
    model.eval()
    return model


def reshape_transform(feats: Union[Tensor, List[Tensor]],
                      max_shape: Tuple[int, int] = (20, 20),
                      is_need_grad: bool = False):
    """Reshape and aggregate feature maps when the input is a multi-layer
    feature map.

    Takes these tensors with different sizes, resizes them to a common shape,
    and concatenates them.
    """
    if len(max_shape) == 1:
        max_shape = max_shape * 2

    if isinstance(feats, torch.Tensor):
        feats = [feats]
    else:
        if is_need_grad:
            raise NotImplementedError('The `grad_base` method does not '
                                      'support output multi-activation layers')

    max_h = max([im.shape[-2] for im in feats])
    max_w = max([im.shape[-1] for im in feats])
    if -1 in max_shape:
        max_shape = (max_h, max_w)
    else:
        max_shape = (min(max_h, max_shape[0]), min(max_w, max_shape[1]))

    activations = []
    for feat in feats:
        activations.append(
            torch.nn.functional.interpolate(
                torch.abs(feat), max_shape, mode='bilinear'))

    activations = torch.cat(activations, axis=1)
    return activations


class BoxAMDetectorWrapper(nn.Module):
    """Wrap the mmdet model class to facilitate handling of non-tensor
    situations during inference."""

    def __init__(self,
                 cfg: ConfigType,
                 checkpoint: str,
                 score_thr: float,
                 device: str = 'cuda:0'):
        super().__init__()
        self.cfg = cfg
        self.device = device
        self.score_thr = score_thr
        self.checkpoint = checkpoint
        self.detector = init_detector(self.cfg, self.checkpoint, device=device)

        pipeline_cfg = copy.deepcopy(self.cfg.test_dataloader.dataset.pipeline)
        pipeline_cfg[0].type = 'mmdet.LoadImageFromNDArray'

        new_test_pipeline = []
        for pipeline in pipeline_cfg:
            if not pipeline['type'].endswith('LoadAnnotations'):
                new_test_pipeline.append(pipeline)
        self.test_pipeline = Compose(new_test_pipeline)

        self.is_need_loss = False
        self.input_data = None
        self.image = None

    def need_loss(self, is_need_loss: bool):
        """Grad-based methods require loss."""
        self.is_need_loss = is_need_loss

    def set_input_data(self,
                       image: np.ndarray,
                       pred_instances: Optional[InstanceData] = None):
        """Set the input data to be used in the next step."""
        self.image = image

        if self.is_need_loss:
            assert pred_instances is not None
            pred_instances = pred_instances.numpy()
            data = dict(
                img=self.image,
                img_id=0,
                gt_bboxes=pred_instances.bboxes,
                gt_bboxes_labels=pred_instances.labels)
            data = self.test_pipeline(data)
        else:
            data = dict(img=self.image, img_id=0)
            data = self.test_pipeline(data)
            data['inputs'] = [data['inputs']]
            data['data_samples'] = [data['data_samples']]
        self.input_data = data

    def __call__(self, *args, **kwargs):
        assert self.input_data is not None
        if self.is_need_loss:
            # Maybe this is a direction that can be optimized
            # self.detector.init_weights()

            self.detector.bbox_head.head_module.training = True
            if hasattr(self.detector.bbox_head, 'featmap_sizes'):
                # Prevent the model algorithm error when calculating loss
                self.detector.bbox_head.featmap_sizes = None

            data_ = {}
            data_['inputs'] = [self.input_data['inputs']]
            data_['data_samples'] = [self.input_data['data_samples']]
            data = self.detector.data_preprocessor(data_, training=False)
            loss = self.detector._run_forward(data, mode='loss')

            if hasattr(self.detector.bbox_head, 'featmap_sizes'):
                self.detector.bbox_head.featmap_sizes = None

            return [loss]
        else:
            self.detector.bbox_head.head_module.training = False
            with torch.no_grad():
                results = self.detector.test_step(self.input_data)
                return results


class BoxAMDetectorVisualizer:
    """Box AM visualization class."""

    def __init__(self,
                 method_class,
                 model: nn.Module,
                 target_layers: List,
                 reshape_transform: Optional[Callable] = None,
                 is_need_grad: bool = False,
                 extra_params: Optional[dict] = None):
        self.target_layers = target_layers
        self.reshape_transform = reshape_transform
        self.is_need_grad = is_need_grad

        if method_class.__name__ == 'AblationCAM':
            batch_size = extra_params.get('batch_size', 1)
            ratio_channels_to_ablate = extra_params.get(
                'ratio_channels_to_ablate', 1.)
            self.cam = AblationCAM(
                model,
                target_layers,
                use_cuda=True if 'cuda' in model.device else False,
                reshape_transform=reshape_transform,
                batch_size=batch_size,
                ablation_layer=extra_params['ablation_layer'],
                ratio_channels_to_ablate=ratio_channels_to_ablate)
        else:
            self.cam = method_class(
                model,
                target_layers,
                use_cuda=True if 'cuda' in model.device else False,
                reshape_transform=reshape_transform,
            )
            if self.is_need_grad:
                self.cam.activations_and_grads.release()

        self.classes = model.detector.dataset_meta['classes']
        self.COLORS = np.random.uniform(0, 255, size=(len(self.classes), 3))

    def switch_activations_and_grads(self, model) -> None:
        """In the grad-based method, we need to switch
        ``ActivationsAndGradients`` layer, otherwise an error will occur."""
        self.cam.model = model

        if self.is_need_grad is True:
            self.cam.activations_and_grads = ActivationsAndGradients(
                model, self.target_layers, self.reshape_transform)
            self.is_need_grad = False
        else:
            self.cam.activations_and_grads.release()
            self.is_need_grad = True

    def __call__(self, img, targets, aug_smooth=False, eigen_smooth=False):
        img = torch.from_numpy(img)[None].permute(0, 3, 1, 2)
        return self.cam(img, targets, aug_smooth, eigen_smooth)[0, :]

    def show_am(self,
                image: np.ndarray,
                pred_instance: InstanceData,
                grayscale_am: np.ndarray,
                with_norm_in_bboxes: bool = False):
        """Normalize the AM to be in the range [0, 1] inside every bounding
        boxes, and zero outside of the bounding boxes."""

        boxes = pred_instance.bboxes
        labels = pred_instance.labels

        if with_norm_in_bboxes is True:
            boxes = boxes.astype(np.int32)
            renormalized_am = np.zeros(grayscale_am.shape, dtype=np.float32)
            images = []
            for x1, y1, x2, y2 in boxes:
                img = renormalized_am * 0
                img[y1:y2, x1:x2] = scale_cam_image(
                    [grayscale_am[y1:y2, x1:x2].copy()])[0]
                images.append(img)

            renormalized_am = np.max(np.float32(images), axis=0)
            renormalized_am = scale_cam_image([renormalized_am])[0]
        else:
            renormalized_am = grayscale_am

        am_image_renormalized = show_cam_on_image(
            image / 255, renormalized_am, use_rgb=False)

        image_with_bounding_boxes = self._draw_boxes(
            boxes, labels, am_image_renormalized, pred_instance.get('scores'))
        return image_with_bounding_boxes

    def _draw_boxes(self,
                    boxes: List,
                    labels: List,
                    image: np.ndarray,
                    scores: Optional[List] = None):
        """draw boxes on image."""
        for i, box in enumerate(boxes):
            label = labels[i]
            color = self.COLORS[label]
            cv2.rectangle(image, (int(box[0]), int(box[1])),
                          (int(box[2]), int(box[3])), color, 2)
            if scores is not None:
                score = scores[i]
                text = str(self.classes[label]) + ': ' + str(
                    round(score * 100, 1))
            else:
                text = self.classes[label]

            cv2.putText(
                image,
                text, (int(box[0]), int(box[1] - 5)),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                color,
                1,
                lineType=cv2.LINE_AA)
        return image


class DetAblationLayer(AblationLayer):
    """Det AblationLayer."""

    def __init__(self):
        super().__init__()
        self.activations = None

    def set_next_batch(self, input_batch_index, activations,
                       num_channels_to_ablate):
        """Extract the next batch member from activations, and repeat it
        num_channels_to_ablate times."""
        if isinstance(activations, torch.Tensor):
            return super().set_next_batch(input_batch_index, activations,
                                          num_channels_to_ablate)

        self.activations = []
        for activation in activations:
            activation = activation[
                input_batch_index, :, :, :].clone().unsqueeze(0)
            self.activations.append(
                activation.repeat(num_channels_to_ablate, 1, 1, 1))

    def __call__(self, x):
        """Go over the activation indices to be ablated, stored in
        self.indices."""
        result = self.activations

        if isinstance(result, torch.Tensor):
            return super().__call__(x)

        channel_cumsum = np.cumsum([r.shape[1] for r in result])
        num_channels_to_ablate = result[0].size(0)  # batch
        for i in range(num_channels_to_ablate):
            pyramid_layer = bisect.bisect_right(channel_cumsum,
                                                self.indices[i])
            if pyramid_layer > 0:
                index_in_pyramid_layer = self.indices[i] - channel_cumsum[
                    pyramid_layer - 1]
            else:
                index_in_pyramid_layer = self.indices[i]
            result[pyramid_layer][i, index_in_pyramid_layer, :, :] = -1000
        return result


class DetBoxScoreTarget:
    """Det Score calculation class.

    In the case of the grad-free method, the calculation method is that
    for every original detected bounding box specified in "bboxes",
    assign a score on how the current bounding boxes match it,

        1. In Bbox IoU
        2. In the classification score.
        3. In Mask IoU if ``segms`` exist.

    If there is not a large enough overlap, or the category changed,
    assign a score of 0. The total score is the sum of all the box scores.

    In the case of the grad-based method, the calculation method is
    the sum of losses after excluding a specific key.
    """

    def __init__(self,
                 pred_instance: InstanceData,
                 match_iou_thr: float = 0.5,
                 device: str = 'cuda:0',
                 ignore_loss_params: Optional[List] = None):
        self.focal_bboxes = pred_instance.bboxes
        self.focal_labels = pred_instance.labels
        self.match_iou_thr = match_iou_thr
        self.device = device
        self.ignore_loss_params = ignore_loss_params
        if ignore_loss_params is not None:
            assert isinstance(self.ignore_loss_params, list)

    def __call__(self, results):
        output = torch.tensor([0.], device=self.device)

        if 'loss_cls' in results:
            # grad-based method
            # results is dict
            for loss_key, loss_value in results.items():
                if 'loss' not in loss_key or \
                        loss_key in self.ignore_loss_params:
                    continue
                if isinstance(loss_value, list):
                    output += sum(loss_value)
                else:
                    output += loss_value
            return output
        else:
            # grad-free method
            # results is DetDataSample
            pred_instances = results.pred_instances
            if len(pred_instances) == 0:
                return output

            pred_bboxes = pred_instances.bboxes
            pred_scores = pred_instances.scores
            pred_labels = pred_instances.labels

            for focal_box, focal_label in zip(self.focal_bboxes,
                                              self.focal_labels):
                ious = torchvision.ops.box_iou(focal_box[None],
                                               pred_bboxes[..., :4])
                index = ious.argmax()
                if ious[0, index] > self.match_iou_thr and pred_labels[
                        index] == focal_label:
                    # TODO: Adaptive adjustment of weights based on algorithms
                    score = ious[0, index] + pred_scores[index]
                    output = output + score
            return output


class SpatialBaseCAM(BaseCAM):
    """CAM that maintains spatial information.

    Gradients are often averaged over the spatial dimension in CAM
    visualization for classification, but this is unreasonable in detection
    tasks. There is no need to average the gradients in the detection task.
    """

    def get_cam_image(self,
                      input_tensor: torch.Tensor,
                      target_layer: torch.nn.Module,
                      targets: List[torch.nn.Module],
                      activations: torch.Tensor,
                      grads: torch.Tensor,
                      eigen_smooth: bool = False) -> np.ndarray:

        weights = self.get_cam_weights(input_tensor, target_layer, targets,
                                       activations, grads)
        weighted_activations = weights * activations
        if eigen_smooth:
            cam = get_2d_projection(weighted_activations)
        else:
            cam = weighted_activations.sum(axis=1)
        return cam


class GradCAM(SpatialBaseCAM, Base_GradCAM):
    """Gradients are no longer averaged over the spatial dimension."""

    def get_cam_weights(self, input_tensor, target_layer, target_category,
                        activations, grads):
        return grads


class GradCAMPlusPlus(SpatialBaseCAM, Base_GradCAMPlusPlus):
    """Gradients are no longer averaged over the spatial dimension."""

    def get_cam_weights(self, input_tensor, target_layers, target_category,
                        activations, grads):
        grads_power_2 = grads**2
        grads_power_3 = grads_power_2 * grads
        # Equation 19 in https://arxiv.org/abs/1710.11063
        sum_activations = np.sum(activations, axis=(2, 3))
        eps = 0.000001
        aij = grads_power_2 / (
            2 * grads_power_2 +
            sum_activations[:, :, None, None] * grads_power_3 + eps)
        # Now bring back the ReLU from eq.7 in the paper,
        # And zero out aijs where the activations are 0
        aij = np.where(grads != 0, aij, 0)

        weights = np.maximum(grads, 0) * aij
        return weights