<gh_stars>0
from ..base_classes import StrictMeta
from abc import abstractmethod
import torch
import logging
logger = logging.getLogger('backboned_unet_attention')

class AttentionModule(metaclass=StrictMeta):
    @abstractmethod
    def compute_attention_map(self, key, query) -> torch.Tensor:
        """
        Compute the attention map between the key and the query
        """

    @abstractmethod
    def forward(self, key, query, value, return_attention_mask):
        """
        Compute the attentive representation fo the value
        """

    def transpose_and_reshape(self, tensor):
        tensor_ = tensor.permute(0, 2, 3, 1)
        tensor_ = tensor_.view(tensor_.size()[0], tensor_.size()[1] * tensor_.size()[2], tensor_.size()[3])
        return tensor_

    def compute_attentive_representation(self, key, query, value):
        """
        Compute the attentive representation of the value in the context of the key and query correlation given
        by the attention map
        :param key: image of dimensions (B, Ck, Hk, Wk)
        :param query: image of dimensions (B, Cq, Hq, Wq)
        :param value: image of dimensions (B, Cv, Hk, Wk)
        :return:
        """
        assert value.shape[0] == key.shape[0]
        assert value.shape[2] == key.shape[2]
        assert value.shape[3] == key.shape[3]

        B, C, H, W = query.shape
        value_ = self.transpose_and_reshape(value)

        attn_mask = self.compute_attention_map(key, query)

        logger.info(f"value_ shape: {value_.shape}")
        result = torch.matmul(attn_mask.permute(0, 2, 1), value_)
        # print(result.)
        result = result.permute(0, 2, 1)  # BATCH x CHANNEL x H x W

        logger.info(f"Attentive representation shape: {result.shape}")
        result = result.view(B, result.shape[1], H, W)
        logger.info(f"Attentive representation final shape: {result.shape}")
        print(f"Attentive representation final shape: {result.shape}")

        return result, attn_mask
from ..base_classes import StrictMeta
from abc import abstractmethod
import torch
import logging
logger = logging.getLogger('backboned_unet_attention')

class AttentionModule(metaclass=StrictMeta):
    @abstractmethod
    def compute_attention_map(self, key, query) -> torch.Tensor:
        """
        Compute the attention map between the key and the query
        """

    @abstractmethod
    def forward(self, key, query, value, return_attention_mask):
        """
        Compute the attentive representation fo the value
        """

    def transpose_and_reshape(self, tensor):
        tensor_ = tensor.permute(0, 2, 3, 1)
        tensor_ = tensor_.view(tensor_.size()[0], tensor_.size()[1] * tensor_.size()[2], tensor_.size()[3])
        return tensor_

    def compute_attentive_representation(self, key, query, value):
        """
        Compute the attentive representation of the value in the context of the key and query correlation given
        by the attention map
        :param key: image of dimensions (B, Ck, Hk, Wk)
        :param query: image of dimensions (B, Cq, Hq, Wq)
        :param value: image of dimensions (B, Cv, Hk, Wk)
        :return:
        """
        assert value.shape[0] == key.shape[0]
        assert value.shape[2] == key.shape[2]
        assert value.shape[3] == key.shape[3]

        B, C, H, W = query.shape
        value_ = self.transpose_and_reshape(value)

        attn_mask = self.compute_attention_map(key, query)

        logger.info(f"value_ shape: {value_.shape}")
        result = torch.matmul(attn_mask.permute(0, 2, 1), value_)
        # print(result.)
        result = result.permute(0, 2, 1)  # BATCH x CHANNEL x H x W

        logger.info(f"Attentive representation shape: {result.shape}")
        result = result.view(B, result.shape[1], H, W)
        logger.info(f"Attentive representation final shape: {result.shape}")
        print(f"Attentive representation final shape: {result.shape}")

        return result, attn_mask
Compute the attention map between the key and the query Compute the attentive representation fo the value Compute the attentive representation of the value in the context of the key and query correlation given         by the attention map         :param key: image of dimensions (B, Ck, Hk, Wk)         :param query: image of dimensions (B, Cq, Hq, Wq)         :param value: image of dimensions (B, Cv, Hk, Wk)         :return: # print(result.) # BATCH x CHANNEL x H x W