Upload folder using huggingface_hub

model/.gitignore

@@ -0,0 +1 @@
+__pycache__/

model/__init__.py

@@ -0,0 +1,90 @@
+from .segmenter import CRIS
+from .segmenter_angular import CRIS_S
+# from .segmenter_ang_nonoise_ddp import CRIS_Wo_Noise
+from loguru import logger
+
+# def build_segmenter(args):
+#     model = CRIS(args)
+#     backbone = []
+#     backbone_no_decay = []
+#     head = []
+#     for k, v in model.named_parameters():
+#         if k.startswith('backbone') and 'positional_embedding' not in k:
+#             backbone.append(v)
+#         elif 'positional_embedding' in k:
+#             backbone_no_decay.append(v)
+#         else:
+#             head.append(v)
+#     print('Backbone with decay: {}, Backbone without decay: {}, Head: {}'.format(
+#         len(backbone), len(backbone_no_decay), len(head)))
+#     param_list = [{
+#         'params': backbone,
+#         'initial_lr': args.lr_multi * args.base_lr
+#     }, {
+#         'params': backbone_no_decay,
+#         'initial_lr': args.lr_multi * args.base_lr,
+#         'weight_decay': 0
+#     }, {
+#         'params': head,
+#         'initial_lr': args.base_lr
+#     }]
+#     return model, param_list
+
+
+def build_segmenter(args):
+    model = CRIS_S(args)
+    backbone = []
+    head = []
+    for k, v in model.named_parameters():
+        if k.startswith('backbone') and 'positional_embedding' not in k:
+            backbone.append(v)
+        else:
+            head.append(v)
+    logger.info('Backbone with decay={}, Head={}'.format(len(backbone), len(head)))
+    param_list = [{
+        'params': backbone,
+        'initial_lr': args.lr_multi * args.base_lr
+    }, {
+        'params': head,
+        'initial_lr': args.base_lr
+    }]
+    return model, param_list
+
+def build_segmenter_original(args):
+    model = CRIS(args)
+    backbone = []
+    head = []
+    for k, v in model.named_parameters():
+        if k.startswith('backbone') and 'positional_embedding' not in k:
+            backbone.append(v)
+        else:
+            head.append(v)
+    logger.info('Backbone with decay={}, Head={}'.format(len(backbone), len(head)))
+    param_list = [{
+        'params': backbone,
+        'initial_lr': args.lr_multi * args.base_lr
+    }, {
+        'params': head,
+        'initial_lr': args.base_lr
+    }]
+    return model, param_list
+
+
+# def build_segmenter_textaug(args):
+#     model = CRIS_Wo_Noise(args)
+#     backbone = []
+#     head = []
+#     for k, v in model.named_parameters():
+#         if k.startswith('backbone') and 'positional_embedding' not in k:
+#             backbone.append(v)
+#         else:
+#             head.append(v)
+#     logger.info('Backbone with decay={}, Head={}'.format(len(backbone), len(head)))
+#     param_list = [{
+#         'params': backbone,
+#         'initial_lr': args.lr_multi * args.base_lr
+#     }, {
+#         'params': head,
+#         'initial_lr': args.base_lr
+#     }]
+#     return model, param_list

model/clip.py

@@ -0,0 +1,554 @@
+from collections import OrderedDict
+from typing import Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1):
+        super().__init__()
+
+        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
+        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+
+        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+
+        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
+
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = None
+        self.stride = stride
+
+        if stride > 1 or inplanes != planes * Bottleneck.expansion:
+            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
+            self.downsample = nn.Sequential(
+                OrderedDict([("-1", nn.AvgPool2d(stride)),
+                             ("0",
+                              nn.Conv2d(inplanes,
+                                        planes * self.expansion,
+                                        1,
+                                        stride=1,
+                                        bias=False)),
+                             ("1", nn.BatchNorm2d(planes * self.expansion))]))
+
+    def forward(self, x: torch.Tensor):
+        identity = x
+
+        out = self.relu(self.bn1(self.conv1(x)))
+        out = self.relu(self.bn2(self.conv2(out)))
+        out = self.avgpool(out)
+        out = self.bn3(self.conv3(out))
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+        return out
+
+
+class AttentionPool2d(nn.Module):
+    def __init__(self,
+                 spacial_dim: int,
+                 embed_dim: int,
+                 num_heads: int,
+                 output_dim: int = None):
+        super().__init__()
+        self.spacial_dim = spacial_dim
+        self.positional_embedding = nn.Parameter(
+            torch.randn(spacial_dim**2 + 1, embed_dim) / embed_dim**0.5)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
+        self.num_heads = num_heads
+        # residual
+        self.connect = nn.Sequential(
+            nn.Conv2d(embed_dim, output_dim, 1, stride=1, bias=False),
+            nn.BatchNorm2d(output_dim))
+
+    def resize_pos_embed(self, pos_embed, input_shpae):
+        """Resize pos_embed weights.
+        Resize pos_embed using bicubic interpolate method.
+        Args:
+            pos_embed (torch.Tensor): Position embedding weights.
+            input_shpae (tuple): Tuple for (downsampled input image height,
+                downsampled input image width).
+            pos_shape (tuple): The resolution of downsampled origin training
+                image.
+            mode (str): Algorithm used for upsampling:
+                ``'nearest'`` | ``'linear'`` | ``'bilinear'`` | ``'bicubic'`` |
+                ``'trilinear'``. Default: ``'nearest'``
+        Return:
+            torch.Tensor: The resized pos_embed of shape [B, C, L_new]
+        """
+        assert pos_embed.ndim == 3, 'shape of pos_embed must be [B, L, C]'
+        pos_h = pos_w = self.spacial_dim
+        cls_token_weight = pos_embed[:, 0]
+        pos_embed_weight = pos_embed[:, (-1 * pos_h * pos_w):]
+        pos_embed_weight = pos_embed_weight.reshape(
+            1, pos_h, pos_w, pos_embed.shape[2]).permute(0, 3, 1, 2)
+        pos_embed_weight = F.interpolate(pos_embed_weight,
+                                         size=input_shpae,
+                                         align_corners=False,
+                                         mode='bicubic')
+        cls_token_weight = cls_token_weight.unsqueeze(1)
+        pos_embed_weight = torch.flatten(pos_embed_weight, 2).transpose(1, 2)
+        # pos_embed = torch.cat((cls_token_weight, pos_embed_weight), dim=1)
+        return pos_embed_weight.transpose(-2, -1)
+
+    def forward(self, x):
+        B, C, H, W = x.size()
+        res = self.connect(x)
+        x = x.reshape(B, C, -1)  # NC(HW)
+        # x = torch.cat([x.mean(dim=-1, keepdim=True), x], dim=-1)  # NC(1+HW)
+        pos_embed = self.positional_embedding.unsqueeze(0)
+        pos_embed = self.resize_pos_embed(pos_embed, (H, W))  # NC(HW)
+        x = x + pos_embed.to(x.dtype)  # NC(HW)
+        x = x.permute(2, 0, 1)  # (HW)NC
+        x, _ = F.multi_head_attention_forward(
+            query=x,
+            key=x,
+            value=x,
+            embed_dim_to_check=x.shape[-1],
+            num_heads=self.num_heads,
+            q_proj_weight=self.q_proj.weight,
+            k_proj_weight=self.k_proj.weight,
+            v_proj_weight=self.v_proj.weight,
+            in_proj_weight=None,
+            in_proj_bias=torch.cat(
+                [self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
+            bias_k=None,
+            bias_v=None,
+            add_zero_attn=False,
+            dropout_p=0,
+            out_proj_weight=self.c_proj.weight,
+            out_proj_bias=self.c_proj.bias,
+            use_separate_proj_weight=True,
+            training=self.training,
+            need_weights=False)
+        x = x.permute(1, 2, 0).reshape(B, -1, H, W)
+        x = x + res
+        x = F.relu(x, True)
+
+        return x
+
+
+class ModifiedResNet(nn.Module):
+    """
+    A ResNet class that is similar to torchvision's but contains the following changes:
+    - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
+    - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
+    - The final pooling layer is a QKV attention instead of an average pool
+    """
+    def __init__(self,
+                 layers,
+                 output_dim,
+                 heads,
+                 input_resolution=224,
+                 width=64):
+        super().__init__()
+        self.output_dim = output_dim
+        self.input_resolution = input_resolution
+
+        # the 3-layer stem
+        self.conv1 = nn.Conv2d(3,
+                               width // 2,
+                               kernel_size=3,
+                               stride=2,
+                               padding=1,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(width // 2)
+        self.conv2 = nn.Conv2d(width // 2,
+                               width // 2,
+                               kernel_size=3,
+                               padding=1,
+                               bias=False)
+        self.bn2 = nn.BatchNorm2d(width // 2)
+        self.conv3 = nn.Conv2d(width // 2,
+                               width,
+                               kernel_size=3,
+                               padding=1,
+                               bias=False)
+        self.bn3 = nn.BatchNorm2d(width)
+        self.avgpool = nn.AvgPool2d(2)
+        self.relu = nn.ReLU(inplace=True)
+
+        # residual layers
+        self._inplanes = width  # this is a *mutable* variable used during construction
+        self.layer1 = self._make_layer(width, layers[0])
+        self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
+        self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
+        self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
+
+        embed_dim = width * 32  # the ResNet feature dimension
+        self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim,
+                                        heads, output_dim)
+
+    def _make_layer(self, planes, blocks, stride=1):
+        layers = [Bottleneck(self._inplanes, planes, stride)]
+
+        self._inplanes = planes * Bottleneck.expansion
+        for _ in range(1, blocks):
+            layers.append(Bottleneck(self._inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        def stem(x):
+            for conv, bn in [(self.conv1, self.bn1), (self.conv2, self.bn2),
+                             (self.conv3, self.bn3)]:
+                x = self.relu(bn(conv(x)))
+            x = self.avgpool(x)
+            return x
+
+        x = x.type(self.conv1.weight.dtype)
+        x = stem(x)
+        x = self.layer1(x)
+        x2 = self.layer2(x)
+        x3 = self.layer3(x2)
+        x4 = self.layer4(x3)
+        x4 = self.attnpool(x4)
+
+        return (x2, x3, x4)
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16."""
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        ret = super().forward(x.type(torch.float32))
+        return ret.type(orig_type)
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self,
+                 d_model: int,
+                 n_head: int,
+                 attn_mask: torch.Tensor = None):
+        super().__init__()
+
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ln_1 = LayerNorm(d_model)
+        self.mlp = nn.Sequential(
+            OrderedDict([("c_fc", nn.Linear(d_model, d_model * 4)),
+                         ("gelu", QuickGELU()),
+                         ("c_proj", nn.Linear(d_model * 4, d_model))]))
+        self.ln_2 = LayerNorm(d_model)
+        self.attn_mask = attn_mask
+
+    def attention(self, x: torch.Tensor):
+        self.attn_mask = self.attn_mask.to(
+            dtype=x.dtype,
+            device=x.device) if self.attn_mask is not None else None
+        return self.attn(x, x, x, need_weights=False,
+                         attn_mask=self.attn_mask)[0]
+
+    def forward(self, x: torch.Tensor):
+        x = x + self.attention(self.ln_1(x))
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+
+class Transformer(nn.Module):
+    def __init__(self,
+                 width: int,
+                 layers: int,
+                 heads: int,
+                 attn_mask: torch.Tensor = None):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.resblocks = nn.Sequential(*[
+            ResidualAttentionBlock(width, heads, attn_mask)
+            for _ in range(layers)
+        ])
+
+    def forward(self, x: torch.Tensor):
+        return self.resblocks(x)
+
+
+class VisionTransformer(nn.Module):
+    def __init__(self, input_resolution: int, patch_size: int, width: int,
+                 layers: int, heads: int, output_dim: int):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.output_dim = output_dim
+        self.conv1 = nn.Conv2d(in_channels=3,
+                               out_channels=width,
+                               kernel_size=patch_size,
+                               stride=patch_size,
+                               bias=False)
+
+        scale = width**-0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(width))
+        self.positional_embedding = nn.Parameter(scale * torch.randn(
+            (input_resolution // patch_size)**2 + 1, width))
+        self.ln_pre = LayerNorm(width)
+
+        self.transformer = Transformer(width, layers, heads)
+
+        self.ln_post = LayerNorm(width)
+        self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
+
+    def forward(self, x: torch.Tensor):
+        x = self.conv1(x)  # shape = [*, width, grid, grid]
+        x = x.reshape(x.shape[0], x.shape[1],
+                      -1)  # shape = [*, width, grid ** 2]
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+        x = torch.cat([
+            self.class_embedding.to(x.dtype) + torch.zeros(
+                x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x
+        ],
+                      dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.positional_embedding.to(x.dtype)
+        x = self.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+
+        # x = self.ln_post(x[:, 0, :])
+        x = self.ln_post(x[:, 1:, :])
+
+        if self.proj is not None:
+            x = x @ self.proj
+
+        return x
+
+
+class CLIP(nn.Module):
+    def __init__(
+            self,
+            embed_dim: int,
+            # vision
+            image_resolution: int,
+            vision_layers: Union[Tuple[int, int, int, int], int],
+            vision_width: int,
+            vision_patch_size: int,
+            # text
+            context_length: int,
+            txt_length: int,
+            vocab_size: int,
+            transformer_width: int,
+            transformer_heads: int,
+            transformer_layers: int):
+        super().__init__()
+
+        self.context_length = context_length
+
+        if isinstance(vision_layers, (tuple, list)):
+            vision_heads = vision_width * 32 // 64
+            self.visual = ModifiedResNet(layers=vision_layers,
+                                         output_dim=embed_dim,
+                                         heads=vision_heads,
+                                         input_resolution=image_resolution,
+                                         width=vision_width)
+        else:
+            vision_heads = vision_width // 64
+            self.visual = VisionTransformer(input_resolution=image_resolution,
+                                            patch_size=vision_patch_size,
+                                            width=vision_width,
+                                            layers=vision_layers,
+                                            heads=vision_heads,
+                                            output_dim=embed_dim)
+
+        self.transformer = Transformer(
+            width=transformer_width,
+            layers=transformer_layers,
+            heads=transformer_heads,
+            attn_mask=self.build_attention_mask(txt_length))
+
+        self.vocab_size = vocab_size
+        self.token_embedding = nn.Embedding(vocab_size, transformer_width)
+        self.positional_embedding = nn.Parameter(
+            torch.empty(self.context_length, transformer_width))
+        self.ln_final = LayerNorm(transformer_width)
+
+        self.text_projection = nn.Parameter(
+            torch.empty(transformer_width, embed_dim))
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        self.token_embedding.requires_grad_ = False
+        self.initialize_parameters()
+
+    def initialize_parameters(self):
+        nn.init.normal_(self.token_embedding.weight, std=0.02)
+        nn.init.normal_(self.positional_embedding, std=0.01)
+
+        if isinstance(self.visual, ModifiedResNet):
+            if self.visual.attnpool is not None:
+                std = self.visual.attnpool.c_proj.in_features**-0.5
+                nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std)
+
+            for resnet_block in [
+                    self.visual.layer1, self.visual.layer2, self.visual.layer3,
+                    self.visual.layer4
+            ]:
+                for name, param in resnet_block.named_parameters():
+                    if name.endswith("bn3.weight"):
+                        nn.init.zeros_(param)
+
+        proj_std = (self.transformer.width**-0.5) * (
+            (2 * self.transformer.layers)**-0.5)
+        attn_std = self.transformer.width**-0.5
+        fc_std = (2 * self.transformer.width)**-0.5
+        for block in self.transformer.resblocks:
+            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+
+        if self.text_projection is not None:
+            nn.init.normal_(self.text_projection,
+                            std=self.transformer.width**-0.5)
+
+    def build_attention_mask(self, context_length):
+        # lazily create causal attention mask, with full attention between the vision tokens
+        # pytorch uses additive attention mask; fill with -inf
+        mask = torch.empty(context_length, context_length)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)  # zero out the lower diagonal
+        return mask
+
+    @property
+    def dtype(self):
+        return self.visual.conv1.weight.dtype
+
+    def encode_image(self, image):
+        return self.visual(image.type(self.dtype))
+
+    def encode_text(self, text):
+        x = self.token_embedding(text).type(
+            self.dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.type(self.dtype)[:x.size(1)]
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x).type(self.dtype)
+
+        # x.shape = [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        state = x[torch.arange(x.shape[0]),
+                  text.argmax(dim=-1)] @ self.text_projection
+        # x = x @ self.text_projection
+        # state = x[torch.arange(x.shape[0]), text.argmax(dim=-1)]
+
+        return x, state
+
+    def forward(self, image, text):
+        image_features = self.encode_image(image)
+        text_features = self.encode_text(text)
+
+        # normalized features
+        image_features = image_features / image_features.norm(dim=-1,
+                                                              keepdim=True)
+        text_features = text_features / text_features.norm(dim=-1,
+                                                           keepdim=True)
+
+        # cosine similarity as logits
+        logit_scale = self.logit_scale.exp()
+        logits_per_image = logit_scale * image_features @ text_features.t()
+        logits_per_text = logits_per_image.t()
+
+        # shape = [global_batch_size, global_batch_size]
+        return logits_per_image, logits_per_text
+
+
+def convert_weights(model: nn.Module):
+    """Convert applicable model parameters to fp16"""
+    def _convert_weights_to_fp16(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.half()
+            if l.bias is not None:
+                l.bias.data = l.bias.data.half()
+
+        if isinstance(l, nn.MultiheadAttention):
+            for attr in [
+                    *[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]],
+                    "in_proj_bias", "bias_k", "bias_v"
+            ]:
+                tensor = getattr(l, attr)
+                if tensor is not None:
+                    tensor.data = tensor.data.half()
+
+        for name in ["text_projection", "proj"]:
+            if hasattr(l, name):
+                attr = getattr(l, name)
+                if attr is not None:
+                    attr.data = attr.data.half()
+
+    model.apply(_convert_weights_to_fp16)
+
+
+def build_model(state_dict: dict, txt_length: int):
+    vit = "visual.proj" in state_dict
+
+    if vit:
+        vision_width = state_dict["visual.conv1.weight"].shape[0]
+        vision_layers = len([
+            k for k in state_dict.keys()
+            if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")
+        ])
+        vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+        grid_size = round(
+            (state_dict["visual.positional_embedding"].shape[0] - 1)**0.5)
+        image_resolution = vision_patch_size * grid_size
+    else:
+        counts: list = [
+            len(
+                set(
+                    k.split(".")[2] for k in state_dict
+                    if k.startswith(f"visual.layer{b}")))
+            for b in [1, 2, 3, 4]
+        ]
+        vision_layers = tuple(counts)
+        vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+        output_width = round(
+            (state_dict["visual.attnpool.positional_embedding"].shape[0] -
+             1)**0.5)
+        vision_patch_size = None
+        assert output_width**2 + 1 == state_dict[
+            "visual.attnpool.positional_embedding"].shape[0]
+        image_resolution = output_width * 32
+
+    embed_dim = state_dict["text_projection"].shape[1]
+    context_length = state_dict["positional_embedding"].shape[0]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(
+        set(
+            k.split(".")[2] for k in state_dict
+            if k.startswith(f"transformer.resblocks")))
+
+    model = CLIP(embed_dim, image_resolution, vision_layers, vision_width,
+                 vision_patch_size, context_length, txt_length, vocab_size,
+                 transformer_width, transformer_heads, transformer_layers)
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        if key in state_dict:
+            del state_dict[key]
+
+    convert_weights(model)
+    model.load_state_dict(state_dict, False)
+    return model.eval()

model/layers.py

@@ -0,0 +1,309 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def conv_layer(in_dim, out_dim, kernel_size=1, padding=0, stride=1):
+    return nn.Sequential(
+        nn.Conv2d(in_dim, out_dim, kernel_size, stride, padding, bias=False),
+        nn.BatchNorm2d(out_dim), nn.ReLU(True))
+
+
+def linear_layer(in_dim, out_dim, bias=False):
+    return nn.Sequential(nn.Linear(in_dim, out_dim, bias),
+                         nn.BatchNorm1d(out_dim), nn.ReLU(True))
+
+
+class CoordConv(nn.Module):
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size=3,
+                 padding=1,
+                 stride=1):
+        super().__init__()
+        self.conv1 = conv_layer(in_channels + 2, out_channels, kernel_size,
+                                padding, stride)
+
+    def add_coord(self, input):
+        b, _, h, w = input.size()
+        x_range = torch.linspace(-1, 1, w, device=input.device)
+        y_range = torch.linspace(-1, 1, h, device=input.device)
+        y, x = torch.meshgrid(y_range, x_range)
+        y = y.expand([b, 1, -1, -1])
+        x = x.expand([b, 1, -1, -1])
+        coord_feat = torch.cat([x, y], 1)
+        input = torch.cat([input, coord_feat], 1)
+        return input
+
+    def forward(self, x):
+        x = self.add_coord(x)
+        x = self.conv1(x)
+        return x
+
+
+class Projector(nn.Module):
+    def __init__(self, word_dim=1024, in_dim=256, kernel_size=3):
+        super().__init__()
+        self.in_dim = in_dim
+        self.kernel_size = kernel_size
+        # visual projector
+        self.vis = nn.Sequential(  # os16 -> os4
+            nn.Upsample(scale_factor=2, mode='bilinear'),
+            conv_layer(in_dim * 2, in_dim * 2, 3, padding=1),
+            nn.Upsample(scale_factor=2, mode='bilinear'),
+            conv_layer(in_dim * 2, in_dim, 3, padding=1),
+            nn.Conv2d(in_dim, in_dim, 1))
+        # textual projector
+        out_dim = 1 * in_dim * kernel_size * kernel_size + 1
+        self.txt = nn.Linear(word_dim, out_dim)
+
+    def forward(self, x, word):
+        '''
+            x: b, 512, 26, 26
+            word: b, 512
+        '''
+        x = self.vis(x)
+        B, C, H, W = x.size()
+        # 1, b*256, 104, 104
+        x = x.reshape(1, B * C, H, W)
+        # txt: b, (256*3*3 + 1) -> b, 256, 3, 3 / b
+        word = self.txt(word)
+        weight, bias = word[:, :-1], word[:, -1]
+        weight = weight.reshape(B, C, self.kernel_size, self.kernel_size)
+        # Conv2d - 1, b*256, 104, 104 -> 1, b, 104, 104
+        out = F.conv2d(x,
+                       weight,
+                       padding=self.kernel_size // 2,
+                       groups=weight.size(0),
+                       bias=bias)
+        out = out.transpose(0, 1)
+        # b, 1, 104, 104
+        return out
+
+
+class TransformerDecoder(nn.Module):
+    def __init__(self,
+                 num_layers,
+                 d_model,
+                 nhead,
+                 dim_ffn,
+                 dropout,
+                 return_intermediate=False):
+        super().__init__()
+        self.layers = nn.ModuleList([
+            TransformerDecoderLayer(d_model=d_model,
+                                    nhead=nhead,
+                                    dim_feedforward=dim_ffn,
+                                    dropout=dropout) for _ in range(num_layers)
+        ])
+        self.num_layers = num_layers
+        self.norm = nn.LayerNorm(d_model)
+        self.return_intermediate = return_intermediate
+
+    @staticmethod
+    def pos1d(d_model, length):
+        """
+        :param d_model: dimension of the model
+        :param length: length of positions
+        :return: length*d_model position matrix
+        """
+        if d_model % 2 != 0:
+            raise ValueError("Cannot use sin/cos positional encoding with "
+                             "odd dim (got dim={:d})".format(d_model))
+        pe = torch.zeros(length, d_model)
+        position = torch.arange(0, length).unsqueeze(1)
+        div_term = torch.exp((torch.arange(0, d_model, 2, dtype=torch.float) *
+                              -(math.log(10000.0) / d_model)))
+        pe[:, 0::2] = torch.sin(position.float() * div_term)
+        pe[:, 1::2] = torch.cos(position.float() * div_term)
+
+        return pe.unsqueeze(1)  # n, 1, 512
+
+    @staticmethod
+    def pos2d(d_model, height, width):
+        """
+        :param d_model: dimension of the model
+        :param height: height of the positions
+        :param width: width of the positions
+        :return: d_model*height*width position matrix
+        """
+        if d_model % 4 != 0:
+            raise ValueError("Cannot use sin/cos positional encoding with "
+                             "odd dimension (got dim={:d})".format(d_model))
+        pe = torch.zeros(d_model, height, width)
+        # Each dimension use half of d_model
+        d_model = int(d_model / 2)
+        div_term = torch.exp(
+            torch.arange(0., d_model, 2) * -(math.log(10000.0) / d_model))
+        pos_w = torch.arange(0., width).unsqueeze(1)
+        pos_h = torch.arange(0., height).unsqueeze(1)
+        pe[0:d_model:2, :, :] = torch.sin(pos_w * div_term).transpose(
+            0, 1).unsqueeze(1).repeat(1, height, 1)
+        pe[1:d_model:2, :, :] = torch.cos(pos_w * div_term).transpose(
+            0, 1).unsqueeze(1).repeat(1, height, 1)
+        pe[d_model::2, :, :] = torch.sin(pos_h * div_term).transpose(
+            0, 1).unsqueeze(2).repeat(1, 1, width)
+        pe[d_model + 1::2, :, :] = torch.cos(pos_h * div_term).transpose(
+            0, 1).unsqueeze(2).repeat(1, 1, width)
+
+        return pe.reshape(-1, 1, height * width).permute(2, 1, 0)  # hw, 1, 512
+
+    def forward(self, vis, txt, pad_mask):
+        '''
+            vis: b, 512, h, w
+            txt: b, L, 512
+            pad_mask: b, L
+        '''
+        B, C, H, W = vis.size()
+        _, L, D = txt.size()
+        # position encoding
+        vis_pos = self.pos2d(C, H, W)
+        txt_pos = self.pos1d(D, L)
+        # reshape & permute
+        vis = vis.reshape(B, C, -1).permute(2, 0, 1)
+        txt = txt.permute(1, 0, 2)
+        # forward
+        output = vis
+        intermediate = []
+        for layer in self.layers:
+            output = layer(output, txt, vis_pos, txt_pos, pad_mask)
+            if self.return_intermediate:
+                # HW, b, 512 -> b, 512, HW
+                intermediate.append(self.norm(output).permute(1, 2, 0))
+
+        if self.norm is not None:
+            # HW, b, 512 -> b, 512, HW
+            output = self.norm(output).permute(1, 2, 0)
+            if self.return_intermediate:
+                intermediate.pop()
+                intermediate.append(output)
+                # [output1, output2, ..., output_n]
+                return intermediate
+            else:
+                # b, 512, HW
+                return output
+        return output
+
+
+class TransformerDecoderLayer(nn.Module):
+    def __init__(self,
+                 d_model=512,
+                 nhead=9,
+                 dim_feedforward=2048,
+                 dropout=0.1):
+        super().__init__()
+        # Normalization Layer
+        self.self_attn_norm = nn.LayerNorm(d_model)
+        self.cross_attn_norm = nn.LayerNorm(d_model)
+        # Attention Layer
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        self.multihead_attn = nn.MultiheadAttention(d_model,
+                                                    nhead,
+                                                    dropout=dropout,
+                                                    kdim=d_model,
+                                                    vdim=d_model)
+        # FFN
+        self.ffn = nn.Sequential(nn.Linear(d_model, dim_feedforward),
+                                 nn.ReLU(True), nn.Dropout(dropout),
+                                 nn.LayerNorm(dim_feedforward),
+                                 nn.Linear(dim_feedforward, d_model))
+        # LayerNorm & Dropout
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+    def with_pos_embed(self, tensor, pos):
+        return tensor if pos is None else tensor + pos.to(tensor.device)
+
+    def forward(self, vis, txt, vis_pos, txt_pos, pad_mask):
+        '''
+            vis: 26*26, b, 512
+            txt: L, b, 512
+            vis_pos: 26*26, 1, 512
+            txt_pos: L, 1, 512
+            pad_mask: b, L
+        '''
+        # Self-Attention
+        vis2 = self.norm1(vis)
+        q = k = self.with_pos_embed(vis2, vis_pos)
+        vis2 = self.self_attn(q, k, value=vis2)[0]
+        vis2 = self.self_attn_norm(vis2)
+        vis = vis + self.dropout1(vis2)
+        # Cross-Attention
+        vis2 = self.norm2(vis)
+        vis2 = self.multihead_attn(query=self.with_pos_embed(vis2, vis_pos),
+                                   key=self.with_pos_embed(txt, txt_pos),
+                                   value=txt,
+                                   key_padding_mask=pad_mask)[0]
+        vis2 = self.cross_attn_norm(vis2)
+        vis = vis + self.dropout2(vis2)
+        # FFN
+        vis2 = self.norm3(vis)
+        vis2 = self.ffn(vis2)
+        vis = vis + self.dropout3(vis2)
+        return vis
+
+
+class FPN(nn.Module):
+    def __init__(self,
+                 in_channels=[512, 1024, 1024],
+                 out_channels=[256, 512, 1024]):
+        super(FPN, self).__init__()
+        # text projection
+        self.txt_proj = linear_layer(in_channels[2], out_channels[2])
+        # fusion 1: v5 & seq -> f_5: b, 1024, 13, 13
+        self.f1_v_proj = conv_layer(in_channels[2], out_channels[2], 1, 0)
+        self.norm_layer = nn.Sequential(nn.BatchNorm2d(out_channels[2]),
+                                        nn.ReLU(True))
+        # fusion 2: v4 & fm -> f_4: b, 512, 26, 26
+        self.f2_v_proj = conv_layer(in_channels[1], out_channels[1], 3, 1)
+        self.f2_cat = conv_layer(out_channels[2] + out_channels[1],
+                                 out_channels[1], 1, 0)
+        # fusion 3: v3 & fm_mid -> f_3: b, 512, 52, 52
+        self.f3_v_proj = conv_layer(in_channels[0], out_channels[0], 3, 1)
+        self.f3_cat = conv_layer(out_channels[0] + out_channels[1],
+                                 out_channels[1], 1, 0)
+        # fusion 4: f_3 & f_4 & f_5 -> fq: b, 256, 26, 26
+        self.f4_proj5 = conv_layer(out_channels[2], out_channels[1], 3, 1)
+        self.f4_proj4 = conv_layer(out_channels[1], out_channels[1], 3, 1)
+        self.f4_proj3 = conv_layer(out_channels[1], out_channels[1], 3, 1)
+        # aggregation
+        self.aggr = conv_layer(3 * out_channels[1], out_channels[1], 1, 0)
+        self.coordconv = nn.Sequential(
+            CoordConv(out_channels[1], out_channels[1], 3, 1),
+            conv_layer(out_channels[1], out_channels[1], 3, 1))
+
+    def forward(self, imgs, state):
+        # v3, v4, v5: 256, 52, 52 / 512, 26, 26 / 1024, 13, 13
+        v3, v4, v5 = imgs
+        # fusion 1: b, 1024, 13, 13
+        # text projection: b, 1024 -> b, 1024
+        state = self.txt_proj(state).unsqueeze(-1).unsqueeze(
+            -1)  # b, 1024, 1, 1
+        f5 = self.f1_v_proj(v5)
+        f5 = self.norm_layer(f5 * state)
+        # fusion 2: b, 512, 26, 26
+        f4 = self.f2_v_proj(v4)
+        f5_ = F.interpolate(f5, scale_factor=2, mode='bilinear')
+        f4 = self.f2_cat(torch.cat([f4, f5_], dim=1))
+        # fusion 3: b, 256, 26, 26
+        f3 = self.f3_v_proj(v3)
+        f3 = F.avg_pool2d(f3, 2, 2)
+        f3 = self.f3_cat(torch.cat([f3, f4], dim=1))
+        # fusion 4: b, 512, 13, 13 / b, 512, 26, 26 / b, 512, 26, 26
+        fq5 = self.f4_proj5(f5)
+        fq4 = self.f4_proj4(f4)
+        fq3 = self.f4_proj3(f3)
+        # query
+        fq5 = F.interpolate(fq5, scale_factor=2, mode='bilinear')
+        fq = torch.cat([fq3, fq4, fq5], dim=1)
+        fq = self.aggr(fq)
+        fq = self.coordconv(fq)
+        # b, 512, 26, 26
+        return fq, f5

model/segmenter.py

@@ -0,0 +1,204 @@
+import random
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from model.clip import build_model
+from .layers import FPN, Projector, TransformerDecoder
+
+
+def MetricLoss(embeddings, n_pos, alpha = 0.5, args = None):
+    # embeddings: ((2*B), C, (H*W))
+    # n_pos : chunk size of positive pairs
+    # args: args
+    # returns: loss
+    metric_loss = 0
+
+    # flatten embeddings
+    B_, C, HW = embeddings.shape
+    emb = torch.mean(embeddings, dim=-1) # (2*B, C)
+    emb_i = emb.unsqueeze(1).repeat(1, B_, 1) # (2*B, 2*B, C)
+    emb_j = emb.unsqueeze(0).repeat(B_, 1, 1) # (2*B, 2*B, C)
+    emb_distance = torch.norm(emb_i - emb_j, dim=-1) # (2*B, 2*B)
+    assert torch.sum(torch.diag(emb_distance[:B_, :B_])) == 0, \
+    "Diagonals are not zero. please check the permutation on the batch"
+    # print("distance metrix : ", emb_distance)
+
+    # positive pairs and loss
+    positive_mask = torch.zeros_like(emb_distance)
+    for i in range(B_//2):
+        positive_mask[2*i, 2*i+1] = 1
+        positive_mask[2*i+1, 2*i] = 1
+    positive_mask.fill_diagonal_(1)
+    positive_loss = torch.sum(emb_distance * positive_mask) / B_
+
+    # negative pairs and loss
+    negative_mask = torch.ones_like(emb_distance) - positive_mask
+
+    if args.div_batch:
+        negative_loss = -1.0 * torch.log(torch.sum(emb_distance * negative_mask) / B_)
+    else:
+        negative_loss = -1.0 * torch.log(torch.sum(emb_distance * negative_mask) / (B_**2 - 2*B_))
+
+    # print(positive_mask, negative_mask)
+
+    metric_loss = alpha * positive_loss + (1-alpha) * negative_loss
+
+    return metric_loss
+
+
+def AngularMetricLoss(embeddings, n_pos, alpha = 0.5, args = None):
+    # embeddings: ((2*B), C, (H*W))
+    # n_pos : chunk size of positive pairs
+    # args: args
+    # returns: loss
+    geometric_loss = 0
+
+    # flatten embeddings
+    B_, C, HW = embeddings.shape
+    emb = torch.mean(embeddings, dim=-1) # (2*B, C)
+    emb_i = emb.unsqueeze(1).repeat(1, B_, 1) # (2*B, 2*B, C)
+    emb_j = emb.unsqueeze(0).repeat(B_, 1, 1) # (2*B, 2*B, C)
+    sim = nn.CosineSimilarity(dim=-1, eps=1e-6)
+    sim_matrix = sim(emb_i, emb_j).reshape(B_, B_) # (2*B , 2*B)
+    print(sim_matrix)
+    assert torch.trace(sim_matrix) == B_, \
+    "similarity diagonals are not one. please check the permutation on the batch"
+    print("similarity metrix : ", sim_matrix)
+    phi = torch.acos(sim_matrix) # (2*B, 2*B)
+    print("phi metrix : ", phi)
+
+    # positive pairs and loss
+    positive_mask = torch.zeros_like(sim_matrix)
+    for i in range(B_//2):
+        positive_mask[2*i, 2*i+1] = 1
+        positive_mask[2*i+1, 2*i] = 1
+    positive_mask.fill_diagonal_(1)
+    positive_loss = torch.sum((phi**2) * positive_mask) / B_
+
+    # negative pairs and loss
+    negative_mask = torch.ones_like(sim_matrix) - positive_mask
+    phi_mask = phi < args.phi_threshold
+    negative_loss = (args.phi_threshold - phi)**2 
+    print(negative_mask * phi_mask)
+    negative_loss = torch.sum(negative_loss * negative_mask * phi_mask) / (B_**2 - 2*B_)
+
+    print("pos loss, neg loss : ", positive_loss, negative_loss)
+
+    geometric_loss = alpha * positive_loss + (1-alpha) * negative_loss
+
+    return geometric_loss
+
+
+class CRIS(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        # Vision & Text Encoder
+        clip_model = torch.jit.load(cfg.clip_pretrain,
+                                    map_location="cpu").eval()
+        self.backbone = build_model(clip_model.state_dict(), cfg.word_len).float()
+        # Multi-Modal FPN
+        self.neck = FPN(in_channels=cfg.fpn_in, out_channels=cfg.fpn_out)
+        # Decoder
+        self.decoder = TransformerDecoder(num_layers=cfg.num_layers,
+                                            d_model=cfg.vis_dim,
+                                            nhead=cfg.num_head,
+                                            dim_ffn=cfg.dim_ffn,
+                                            dropout=cfg.dropout,
+                                            return_intermediate=cfg.intermediate)
+        # Projector
+        self.proj = Projector(cfg.word_dim, cfg.vis_dim // 2, 3)
+        self.metric_learning = cfg.metric_learning
+        self.positive_strength = cfg.positive_strength
+        self.metric_loss_weight = cfg.metric_loss_weight
+        self.eps = cfg.ptb_rate
+        self.cfg = cfg
+
+    def forward(self, image, text, target=None):
+        '''
+            img: b, 3, h, w
+            word: b, words
+            word_mask: b, words
+            if self.metric_learning:
+                word: b, 2, words
+                word_mask: b, 2, words
+            mask: b, 1, h, w
+        '''
+        metric_learning_flag = (self.metric_learning and self.training)
+        metric_loss = 0
+
+        # 1.Resizing : if metric learning, batch size of the word is doubled
+        if metric_learning_flag:
+            #print("image shape : ", image.shape)
+            b, c, h, w = image.size()
+            # duplicate image and segmentation mask
+            if image is not None:
+                image = torch.cat([image, image], dim=0)
+                image = image.reshape(-1, b, c, h, w).transpose(0, 1).reshape(-1, c, h, w)
+            if target is not None:
+                target = torch.cat([target, target], dim=0)
+                target = target.reshape(-1, b, 1, h, w).transpose(0, 1).reshape(-1, 1, h, w)
+            # duplicate noise mask
+            b_, n_, l_ = text.size()
+            assert n_ == 2 ,"word size should be 2"
+            noise_mask = (text[:, 0, :] == text[:, 1, :])
+            noise_mask = torch.all(noise_mask, dim=-1)
+            noise_mask = noise_mask.unsqueeze(-1).repeat(1, 2).reshape(-1) # 2*b_
+            assert noise_mask.shape[0] == b_ * 2, "noise mask shape should be 2*b_"
+            text = text.reshape(b_ * 2, l_) # 2*b, l
+
+        # print("text shape : ", text.shape)
+        # print("image shape : ", image.shape)
+        # print("target shape : ", target.shape)
+        # print(torch.sum(image[0::2]) == torch.sum(image[1::2]))
+        # print(torch.sum(target[0::2]) == torch.sum(target[1::2]))
+        
+        # padding mask used in decoder
+        pad_mask = torch.zeros_like(text).masked_fill_(text == 0, 1).bool()
+        # vis: C3 / C4 / C5
+        # word: b, length, 1024
+        # state: b, 1024
+        vis = self.backbone.encode_image(image)
+        word, state = self.backbone.encode_text(text)
+
+        b_, d_ = state.size()
+        assert b_ == word.size(0), "batch size of state and word should be same"
+
+
+        # 2. State Noising Step : if number of caption is 1,
+        # add noise to the corresponding indices
+        if metric_learning_flag :
+            noise = torch.randn_like(state) * self.eps
+            state[noise_mask] = state[noise_mask] + noise[noise_mask]
+
+        # print("shape of word, state : ", word.shape, state.shape)
+
+        # b, 512, 26, 26 (C4)
+        a3, a4, a5 = vis
+        # print("vis shape in model " , a3.shape, a4.shape, a5.shape)
+        fq, f5 = self.neck(vis, state)
+        b, c, h, w = fq.size()
+        fq = self.decoder(fq, word, pad_mask)
+        # print("decoder output shape : ", fq.shape)
+        # 3. Get metric loss
+        if metric_learning_flag:
+            metric_loss = MetricLoss(fq, 2, alpha=self.positive_strength, args = self.cfg)
+            
+        fq = fq.reshape(b, c, h, w)
+
+        # b, 1, 104, 104
+        pred = self.proj(fq, state)
+
+        if self.training:
+            # resize mask
+            if pred.shape[-2:] != target.shape[-2:]:
+                target = F.interpolate(target, pred.shape[-2:],
+                                    mode='nearest').detach()
+            loss = F.binary_cross_entropy_with_logits(pred, target)
+            # 4. if metric learning, add metric loss and normalize
+            if metric_learning_flag:
+                #print("CE loss : ", loss, "metric loss : ", metric_loss)
+                loss = (loss + self.metric_loss_weight * metric_loss) / (1+self.metric_loss_weight)
+            return pred.detach(), target, loss
+        else:
+            return pred.detach()

model/segmenter_angular.py

@@ -0,0 +1,163 @@
+import random
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from model.clip import build_model
+
+from .layers import FPN, Projector, TransformerDecoder
+
+
+# def MetricLoss(embeddings, n_pos, alpha = 0.5, args = None):
+#     # embeddings: ((2*B), C, (H*W))
+#     # n_pos : chunk size of positive pairs
+#     # args: args
+#     # returns: loss
+#     metric_loss = 0
+
+#     # flatten embeddings
+#     B_, C, HW = embeddings.shape
+#     emb = torch.mean(embeddings, dim=-1) # (2*B, C)
+#     emb_i = emb.unsqueeze(1).repeat(1, B_, 1) # (2*B, 2*B, C)
+#     emb_j = emb.unsqueeze(0).repeat(B_, 1, 1) # (2*B, 2*B, C)
+#     emb_distance = torch.norm(emb_i - emb_j, dim=-1) # (2*B, 2*B)
+#     assert torch.sum(torch.diag(emb_distance[:B_, :B_])) == 0, \
+#     "Diagonals are not zero. please check the permutation on the batch"
+#     # print("distance metrix : ", emb_distance)
+
+#     # positive pairs and loss
+#     positive_mask = torch.zeros_like(emb_distance)
+#     for i in range(B_//2):
+#         positive_mask[2*i, 2*i+1] = 1
+#         positive_mask[2*i+1, 2*i] = 1
+#     positive_mask.fill_diagonal_(1)
+#     positive_loss = torch.sum(emb_distance * positive_mask) / B_
+
+#     # negative pairs and loss
+#     negative_mask = torch.ones_like(emb_distance) - positive_mask
+#     negative_loss = -1.0 * torch.log(torch.sum(emb_distance * negative_mask) / (B_**2 - 2*B_))
+
+#     # print(positive_mask, negative_mask)
+
+#     metric_loss = alpha * positive_loss + (1-alpha) * negative_loss
+
+#     return metric_loss
+
+
+
+class CRIS_S(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        # Vision & Text Encoder
+        clip_model = torch.jit.load(cfg.clip_pretrain,
+                                    map_location="cpu").eval()
+        self.backbone = build_model(clip_model.state_dict(), cfg.word_len).float()
+        # Multi-Modal FPN
+        self.neck = FPN(in_channels=cfg.fpn_in, out_channels=cfg.fpn_out)
+        # Decoder
+        self.decoder = TransformerDecoder(num_layers=cfg.num_layers,
+                                            d_model=cfg.vis_dim,
+                                            nhead=cfg.num_head,
+                                            dim_ffn=cfg.dim_ffn,
+                                            dropout=cfg.dropout,
+                                            return_intermediate=cfg.intermediate)
+        # Projector
+        self.proj = Projector(cfg.word_dim, cfg.vis_dim // 2, 3)
+        self.metric_learning = cfg.metric_learning
+        self.positive_strength = cfg.positive_strength
+        self.metric_loss_weight = cfg.metric_loss_weight
+        self.eps = cfg.ptb_rate
+        self.cfg = cfg
+
+    def forward(self, image, text, target=None):
+        '''
+            img: b, 3, h, w
+            word: b, words
+            word_mask: b, words
+            if self.metric_learning:
+                word: b, 2, words
+                word_mask: b, 2, words
+            mask: b, 1, h, w
+        '''
+        metric_learning_flag = (self.metric_learning and self.training)
+        # TODO : mixing option btw distance & angular loss
+        mix_distance_angular = False
+        metric_loss = 0
+
+        # 1.Resizing : if metric learning, batch size of the word is doubled
+        if metric_learning_flag:
+            #print("image shape : ", image.shape)
+            b, c, h, w = image.size()
+            # duplicate image and segmentation mask
+            if image is not None:
+                image = torch.cat([image, image], dim=0)
+                image = image.reshape(-1, b, c, h, w).transpose(0, 1).reshape(-1, c, h, w)
+            if target is not None:
+                target = torch.cat([target, target], dim=0)
+                target = target.reshape(-1, b, 1, h, w).transpose(0, 1).reshape(-1, 1, h, w)
+            # duplicate noise mask
+            b_, n_, l_ = text.size()
+            assert n_ == 2 ,"word size should be 2"
+            noise_mask = (text[:, 0, :] == text[:, 1, :])
+            noise_mask = torch.all(noise_mask, dim=-1)
+            noise_mask = noise_mask.unsqueeze(-1).repeat(1, 2).reshape(-1) # 2*b_
+            assert noise_mask.shape[0] == b_ * 2, "noise mask shape should be 2*b_"
+            text = text.reshape(b_ * 2, l_) # 2*b, l
+
+        # print("text shape : ", text.shape)
+        # print("image shape : ", image.shape)
+        # print("target shape : ", target.shape)
+        # print(torch.sum(image[0::2]) == torch.sum(image[1::2]))
+        # print(torch.sum(target[0::2]) == torch.sum(target[1::2]))
+        
+        # padding mask used in decoder
+        pad_mask = torch.zeros_like(text).masked_fill_(text == 0, 1).bool()
+        # vis: C3 / C4 / C5
+        # word: b, length, 1024
+        # state: b, 1024
+        vis = self.backbone.encode_image(image)
+        word, state = self.backbone.encode_text(text)
+
+        b_, d_ = state.size()
+        assert b_ == word.size(0), "batch size of state and word should be same"
+
+
+        # 2. State Noising Step : if number of caption is 1,
+        # add noise to the corresponding indices
+        if metric_learning_flag :
+            noise = torch.randn_like(state) * self.eps
+            state[noise_mask] = state[noise_mask] + noise[noise_mask]
+
+
+        # b, 512, 26, 26 (C4)
+        a3, a4, a5 = vis
+        fq, f5 = self.neck(vis, state)
+        b, c, h, w = fq.size()
+        fq = self.decoder(fq, word, pad_mask)
+        metric_tensor = fq
+        # if metric_learning_flag:
+        #     metric_loss = AngularMetricLoss(fq, 2, alpha=self.positive_strength, args = self.cfg) #  (1-self.positive_strength) *
+        #     if mix_distance_angular:
+        #         metric_loss += MetricLoss(fq, 2, alpha=self.positive_strength, args = self.cfg) #  self.positive_strength * 
+        fq = fq.reshape(b, c, h, w)
+
+        # b, 1, 104, 104
+        pred = self.proj(fq, state)
+
+        if self.training:
+            # resize mask
+            if pred.shape[-2:] != target.shape[-2:]:
+                target = F.interpolate(target, pred.shape[-2:],
+                                    mode='nearest').detach()
+            CE_loss = F.binary_cross_entropy_with_logits(pred, target)
+
+            # 4. if metric learning, add metric loss and normalize
+            # if metric_learning_flag:
+            #     loss = (loss + self.metric_loss_weight * metric_loss) / (1+self.metric_loss_weight)
+            #     safety_loss = loss * 0.
+            #     loss = loss + safety_loss
+
+            return pred.detach(), target, CE_loss, metric_tensor
+        else:
+            #print(self.cfg.gpu, f"; loss = {loss}")
+            return pred.detach()

model/segmenter_verbonly.py

@@ -0,0 +1,375 @@
+import random
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from model.clip import build_model
+
+from .layers import FPN, Projector, TransformerDecoder
+
+class CRIS_VerbOnly(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        # Vision & Text Encoder
+        clip_model = torch.jit.load(cfg.clip_pretrain,
+                                    map_location="cpu").eval()
+        self.backbone = build_model(clip_model.state_dict(), cfg.word_len).float()
+        # Multi-Modal FPN
+        self.neck = FPN(in_channels=cfg.fpn_in, out_channels=cfg.fpn_out)
+        # Decoder
+        self.decoder = TransformerDecoder(num_layers=cfg.num_layers,
+                                            d_model=cfg.vis_dim,
+                                            nhead=cfg.num_head,
+                                            dim_ffn=cfg.dim_ffn,
+                                            dropout=cfg.dropout,
+                                            return_intermediate=cfg.intermediate)
+        # Projector
+        self.proj = Projector(cfg.word_dim, cfg.vis_dim // 2, 3)
+        self.metric_learning = False # cfg.metric_learning
+        self.positive_strength = cfg.positive_strength
+        self.metric_loss_weight = cfg.metric_loss_weight
+        self.eps = cfg.ptb_rate
+        self.cfg = cfg
+
+
+
+
+    def forward(self, image, text, target=None, verb=None):
+        '''
+            image: b, 3, h, w
+            text: b, words
+            target: b, 1, h, w
+            verb: b, words (if applicable, only used in training mode for contrastive learning)
+        '''
+
+        sentences, images, targets, pad_masks = [], [], [], []
+        
+        if self.training: 
+            verb_masks = [] 
+            cl_masks = []
+
+            for idx in range(len(text)):
+                sentences.append(text[idx])
+                images.append(image[idx])
+                targets.append(target[idx])
+                pad_masks.append(torch.zeros_like(text[idx]).masked_fill_(text[idx] == 0, 1).bool())
+
+                # If verb exists, process it
+                if verb[idx].numel() > 0 and verb[idx].sum().item() > 0:
+                    verb_masks.extend([1, 1])  # Both original sentence and verb are marked
+                    cl_masks.extend([0, 1])    # Only verb gets marked for exclusion from CE loss
+                    sentences.append(verb[idx])
+                    images.append(image[idx])
+                    targets.append(target[idx])
+                    pad_masks.append(torch.zeros_like(verb[idx]).masked_fill_(verb[idx] == 0, 1).bool())
+                else:
+                    verb_masks.append(0)
+                    cl_masks.append(0)
+
+
+            sentences = torch.stack(sentences)
+            images = torch.stack(images)
+            targets = torch.stack(targets)
+            pad_masks = torch.stack(pad_masks)
+            verb_masks = torch.tensor(verb_masks, dtype=torch.bool)
+            cl_masks = torch.tensor(cl_masks, dtype=torch.bool)  
+
+        else: 
+            sentences = text
+            images = image
+            targets = target
+            pad_masks = torch.zeros_like(text).masked_fill_(text == 0, 1).bool()
+
+        # Encoding images and text
+        vis = self.backbone.encode_image(images)
+        word, state = self.backbone.encode_text(sentences)
+
+        # FPN neck and decoder
+        fq, f5 = self.neck(vis, state)
+        b, c, h, w = fq.size()
+        fq = self.decoder(fq, word, pad_masks)
+        metric_tensor = fq  # b, c, h*w
+        fq = fq.reshape(b, c, h, w)
+
+        # Final prediction
+        pred = self.proj(fq, state)
+
+        if self.training:
+            if pred.shape[-2:] != targets.shape[-2:]:
+                targets = F.interpolate(targets, pred.shape[-2:], mode='nearest').detach()
+
+            loss = F.binary_cross_entropy_with_logits(pred[~cl_masks], targets[~cl_masks])
+
+            if self.metric_learning:
+                metric_loss = self.compute_metric_loss(metric_tensor, verb_masks, args=self.cfg)
+                loss = (loss + self.metric_loss_weight * metric_loss) / (1 + self.metric_loss_weight)
+            
+            return pred.detach(), targets, loss
+
+        return pred.detach()  # In eval mode, only return the predictions
+
+
+
+
+    def return_mask_hponly(self, emb_distance, verb_mask=None):
+        B_, B_ = emb_distance.shape
+        positive_mask = torch.zeros_like(emb_distance)
+        positive_mask.fill_diagonal_(1)  # Set diagonal elements to 1 for all cases
+        
+        if B_ < len(verb_mask):
+            # If B_ equals to 2*K (double the number of verb phrase)
+            for i in range(B_ // 2):
+                positive_mask[2 * i, 2 * i + 1] = 1
+                positive_mask[2 * i + 1, 2 * i] = 1
+        else:
+            # Process the case where we have a mix of sentences with and without verbs
+            i = 0
+            while i < B_:
+                if verb_mask[i] == 1:
+                    positive_mask[i, i + 1] = 1
+                    positive_mask[i + 1, i] = 1
+                    i += 2
+                else:
+                    i += 1  
+        negative_mask = torch.ones_like(emb_distance) - positive_mask
+        return positive_mask, negative_mask
+
+
+    def return_mask_hphn(self, emb_distance, positive_verbs, negative_verbs, verb_mask):
+        B_, B_ = emb_distance.shape
+        positive_mask = torch.zeros_like(emb_distance)
+        negative_mask = torch.ones_like(emb_distance)
+        positive_mask.fill_diagonal_(1)
+
+        if B_ < len(verb_mask):
+            # Considering only verbs that pass the verb_mask filter
+            positive_verbs = torch.tensor(positive_verbs)[verb_mask]
+            negative_verbs = torch.tensor(negative_verbs)[verb_mask]
+
+            # Exclude hard negatives from both masks (diagonal)
+            for i in range(B_):
+                if negative_verbs[i] == 1:
+                    positive_mask[i, i] = 0
+                    negative_mask[i, i] = 0
+
+            i = 0
+            while i < B_:
+                if positive_verbs[i] == 1:  
+                    if i + 1 < B_ and positive_verbs[i + 1] == 1:
+                        positive_mask[i, i + 1] = 1
+                        positive_mask[i + 1, i] = 1
+                    i += 2 
+                else:
+                    i += 1
+        else:
+            # Exclude hard negatives from both masks (diagonal)
+            for i in range(B_):
+                if negative_verbs[i] == 1:
+                    positive_mask[i, i] = 0
+                    negative_mask[i, i] = 0
+
+            # Apply the positive pairs logic similarly as above
+            i = 0
+            while i < B_:
+                if positive_verbs[i] == 1 and i + 1 < B_ and positive_verbs[i + 1] == 1:
+                    positive_mask[i, i + 1] = 1
+                    positive_mask[i + 1, i] = 1
+                    i += 2
+                else:
+                    i += 1
+
+        negative_mask = negative_mask - positive_mask
+
+        return positive_mask, negative_mask
+
+
+
+
+
+    def compute_contrastive_loss(self, fq, state, verb_masks, temperature=0.05):
+        """
+        Compute contrastive loss (NCE) only for the samples with verb phrases.
+        fq: shape (b, c, h*w) -> Encoded image features
+        state: shape (b, d) -> Encoded text features (word representations)
+        verb_masks: boolean mask indicating samples containing verb phrases
+        temperature: scaling factor for contrastive loss
+        """
+
+        # Extract only the samples that contain verbs using verb_masks
+        fq_verb_samples = fq[verb_masks]  # (num_verbs, c, h*w)
+        state_verb_samples = state[verb_masks]  # (num_verbs, d)
+
+        fq_verb_samples = F.normalize(fq_verb_samples, p=2, dim=1)  
+        state_verb_samples = F.normalize(state_verb_samples, p=2, dim=1) 
+        
+        # Compute the inner product between language conditioned feature output and encoded text (verb phrases)
+        fq_verb_flat = fq_verb_samples.view(fq_verb_samples.size(0), -1) 
+        logits = torch.matmul(fq_verb_flat, state_verb_samples.t())  
+        logits = logits / temperature  
+
+        # Create labels for the contrastive loss (positive pairs are diagonals)
+        labels = torch.arange(logits.size(0), device=logits.device)
+
+        contrastive_loss = F.cross_entropy(logits, labels)
+        
+        return contrastive_loss
+
+
+    # cosine sim only on metric_tensor
+    def AngularContrastiveLoss_1(self, embeddings, verb_mask, alpha=0.5, m=0.5, tau=0.05, args=None):
+        """
+        Angular Margin Contrastive Loss function.
+        - \( \theta_{i, i^*} \) represents the cosine similarity between the anchor \( h_i \) and the positive sample \( h_{i^*} \).
+        - An angular margin \( m \) is added to increase the distance between the positive and negative pairs.
+        - \( \tau \) is a temperature scaling factor to control the sharpness of the probability distribution.
+
+        https://aclanthology.org/2022.acl-long.336.pdf 
+        \[
+        \mathcal{L}_{arc} = -\log \frac{\exp\left(\cos(\theta_{i,i^*} + m)/\tau\right)}{\exp\left(\cos(\theta_{i,i^*} + m)/\tau\right) + \sum_{j \neq i} \exp\left(\cos(\theta_{i,j})/\tau\right)}
+        \]
+
+        Args:
+            embeddings: Encoded embeddings with shape (B, C, H*W) for image-text fused features.
+            verb_mask: A mask indicating the samples with verb phrases.
+            alpha: Weight for balancing positive and negative loss components.
+            m: Angular margin to add to the cosine similarity of positive pairs.
+            tau: Temperature scaling factor for softmax.
+            args: Optional arguments for additional control.
+
+        Returns:
+            geometric_loss: Calculated Angular Metric Loss.
+        """
+
+        # Get batch size and feature dimensions
+        B_, C, HW = embeddings.shape
+
+        # Mean pooling across the spatial dimension (H*W) and normalize embeddings
+        emb = torch.mean(embeddings[verb_mask], dim=-1)  # (B_, C)
+        emb = F.normalize(emb, p=2, dim=1)  # Normalize the embeddings
+
+        # Create cosine similarity matrix
+        sim = nn.CosineSimilarity(dim=-1, eps=1e-6)
+
+        # Pairwise cosine similarities
+        emb_i = emb.unsqueeze(1).repeat(1, B_, 1)  # Expand emb_i to pair with all other embeddings
+        emb_j = emb.unsqueeze(0).repeat(B_, 1, 1)  # Expand emb_j to pair with all other embeddings
+        sim_matrix = sim(emb_i, emb_j).reshape(B_, B_)  # (B_, B_)
+
+        # Clamp values to avoid numerical instability
+        sim_matrix = torch.clamp(sim_matrix, min=-0.9999, max=0.9999)
+
+        # Apply angular margin for positive pairs
+        positive_mask = torch.eye(B_, device=embeddings.device).bool()  # Diagonal is the positive pairs
+        sim_matrix_with_margin = sim_matrix.clone()
+
+        # Apply the angular margin `m` only to positive pairs (diagonal)
+        sim_matrix_with_margin[positive_mask] = torch.cos(torch.acos(sim_matrix[positive_mask]) + m)
+
+        # Scale logits with temperature
+        logits = sim_matrix_with_margin / tau
+
+        # Compute the softmax loss for all pairs
+        exp_logits = torch.exp(logits)
+        pos_exp_logits = exp_logits[positive_mask]
+        total_exp_logits = exp_logits.sum(dim=-1)
+
+        # Compute the final loss: L_arc = -log(e^(cos(theta + m)/tau) / sum(e^(cos(theta)/tau)))
+        positive_loss = -torch.log(pos_exp_logits / total_exp_logits)
+
+        # Average the loss over the batch size
+        angular_loss = positive_loss.mean()
+
+        return angular_loss
+
+
+    # cosine similarity on metric_tensor (image-text) and state (text eos) 
+    def AngularContrastiveLoss_2(self, fq, state, verb_masks, alpha=1.0, margin=0.5, temperature=0.05):
+        """
+        Angular Margin Contrastive Loss function.
+        - \( \theta_{i, i^*} \) represents the cosine similarity between the anchor \( h_i \) and the positive sample \( h_{i^*} \).
+        - An angular margin \( m \) is added to increase the distance between the positive and negative pairs.
+        - \( \tau \) is a temperature scaling factor to control the sharpness of the probability distribution.
+
+        https://aclanthology.org/2022.acl-long.336.pdf 
+        \[
+        \mathcal{L}_{arc} = -\log \frac{\exp\left(\cos(\theta_{i,i^*} + m)/\tau\right)}{\exp\left(\cos(\theta_{i,i^*} + m)/\tau\right) + \sum_{j \neq i} \exp\left(\cos(\theta_{i,j})/\tau\right)}
+        \]
+
+
+        fq: (b, c, h*w) -> Encoded language-fused multimodal feature (metric_tensor) 
+        state: (b, d) -> Encoded text features (word representations)
+        verb_masks: boolean mask indicating samples containing verb phrases
+        alpha: weight for positive samples
+        margin: the angular margin to enforce between positive pairs
+        temperature: scaling factor for contrastive loss
+        """
+
+        # Select only the verb-containing samples
+        # Assume c equals to d (CLIP model backbone)
+        fq_verb_samples = torch.mean(fq[verb_masks], dim=-1)  # (num_verbs, d)
+        state_verb_samples = state[verb_masks]  # (num_verbs, d)
+        fq_verb_samples = F.normalize(fq_verb_samples, p=2, dim=1)  # (num_verbs, d)
+        state_verb_samples = F.normalize(state_verb_samples, p=2, dim=1)  # (num_verbs, d)
+
+        # Compute cosine similarity (logits) between image and text features
+        logits = torch.matmul(fq_verb_samples, state_verb_samples.t())  # (num_verbs, num_verbs)
+
+        # Apply the angular margin to positive pairs (diagonal entries)
+        diagonal_indices = torch.arange(logits.size(0), device=logits.device)
+        positive_logits = logits[diagonal_indices, diagonal_indices]
+        positive_logits_with_margin = positive_logits + margin
+
+        # Replace the diagonal (positive) entries with the margin-added values
+        logits[diagonal_indices, diagonal_indices] = positive_logits_with_margin
+        logits = logits / temperature
+
+        # Create positive mask (diagonal) and negative mask (non-diagonal)
+        positive_mask = torch.eye(logits.size(0), device=logits.device).bool()  # Diagonal for positive pairs
+        negative_mask = ~positive_mask  # Non-diagonal for negative pairs
+        exp_logits = torch.exp(logits)  # Exponentials of logits
+
+        # Positive and negative softmax components
+        pos_exp_logits = exp_logits[positive_mask].view(-1)  # Positive pairs (diagonal entries)
+        neg_exp_logits = exp_logits[negative_mask].view(logits.size(0), -1).sum(dim=1)  # Sum of negative pairs
+
+        # Final loss: -log(e^(cos(theta + m)/tau) / (e^(cos(theta + m)/tau) + sum(e^(cos(theta)/tau)))
+        positive_loss = -torch.log(pos_exp_logits / (pos_exp_logits + neg_exp_logits))
+
+        loss = positive_loss.mean()
+
+        return loss
+
+
+    # def AngularMetricLoss_Seunghoon(self, embeddings, n_pos , verb_mask,  alpha = 0.5, args = None):
+    #     # embeddings: ((2*B), C, (H*W))
+    #     # n_pos : chunk size of positive pairs
+    #     # args: args
+    #     # returns: loss
+    #     geometric_loss = 0
+
+    #     # flatten embeddings
+    #     B_, C, HW = embeddings.shape
+    #     emb = torch.mean(embeddings[verb_mask], dim=-1) 
+    #     emb_i = emb.unsqueeze(1).repeat(1, B_, 1) 
+    #     emb_j = emb.unsqueeze(0).repeat(B_, 1, 1) 
+    #     sim = nn.CosineSimilarity(dim=-1, eps=1e-6)
+    #     sim_matrix = sim(emb_i, emb_j).reshape(B_, B_) 
+    #     sim_matrix = torch.clamp(sim_matrix, min=-0.9999, max=0.9999)
+    #     phi = torch.acos(sim_matrix) 
+
+    #     # positive pairs and loss
+    #     positive_mask = torch.zeros_like(sim_matrix)
+    #     positive_mask.fill_diagonal_(1)
+    #     positive_loss = torch.sum((phi**2) * positive_mask) / B_
+
+    #     # negative pairs and loss
+    #     negative_mask = torch.ones_like(sim_matrix) - positive_mask
+    #     phi_mask = phi < args.phi_threshold
+    #     negative_loss = (args.phi_threshold - phi)**2 
+    #     #print(negative_mask * phi_mask)
+    #     negative_loss = torch.sum(negative_loss * negative_mask * phi_mask) / (B_**2 - 2*B_)
+
+    #     #print("pos loss, neg loss : ", positive_loss, negative_loss)
+    #     geometric_loss = alpha * positive_loss + (1-alpha) * negative_loss
+
+    #     return geometric_loss

model/segmenter_verbonly_hardneg.py

@@ -0,0 +1,204 @@
+import random
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from model.clip import build_model
+
+from .layers import FPN, Projector, TransformerDecoder
+
+class CRIS_VerbOnly(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        # Vision & Text Encoder
+        clip_model = torch.jit.load(cfg.clip_pretrain,
+                                    map_location="cpu").eval()
+        self.backbone = build_model(clip_model.state_dict(), cfg.word_len).float()
+        # Multi-Modal FPN
+        self.neck = FPN(in_channels=cfg.fpn_in, out_channels=cfg.fpn_out)
+        # Decoder
+        self.decoder = TransformerDecoder(num_layers=cfg.num_layers,
+                                            d_model=cfg.vis_dim,
+                                            nhead=cfg.num_head,
+                                            dim_ffn=cfg.dim_ffn,
+                                            dropout=cfg.dropout,
+                                            return_intermediate=cfg.intermediate)
+        # Projector
+        self.proj = Projector(cfg.word_dim, cfg.vis_dim // 2, 3)
+        self.metric_learning = False # cfg.metric_learning
+        self.positive_strength = cfg.positive_strength
+        self.metric_loss_weight = cfg.metric_loss_weight
+        self.eps = cfg.ptb_rate
+        self.cfg = cfg
+
+
+    def forward(self, image, text, target=None, hardpos=None, hardneg=None):
+        '''
+            image: b, 3, h, w
+            text: b, words
+            target: b, 1, h, w
+            verb: b, words (if applicable, only used in training mode for contrastive learning)
+        '''
+
+        sentences, images, targets, pad_masks = [], [], [], []
+        positive_verbs, negative_verbs = [], []
+        
+        
+        if self.training: 
+            verb_masks = [] 
+            cl_masks = []
+
+            for idx in range(len(text)):
+                sentences.append(text[idx])
+                images.append(image[idx])
+                targets.append(target[idx])
+                pad_masks.append(torch.zeros_like(text[idx]).masked_fill_(text[idx] == 0, 1).bool())
+
+
+
+                # If verb exists, process it
+                if verb[idx].numel() > 0 and verb[idx].sum().item() > 0:
+                    verb_masks.extend([1, 1])  # Both original sentence and verb are marked
+                    cl_masks.extend([0, 1])    # Only verb gets marked for exclusion from CE loss
+                    sentences.append(verb[idx])
+                    images.append(image[idx])
+                    targets.append(target[idx])
+                    pad_masks.append(torch.zeros_like(verb[idx]).masked_fill_(verb[idx] == 0, 1).bool())
+                else:
+                    verb_masks.append(0)
+                    cl_masks.append(0)
+
+
+            sentences = torch.stack(sentences)
+            images = torch.stack(images)
+            targets = torch.stack(targets)
+            pad_masks = torch.stack(pad_masks)
+            verb_masks = torch.tensor(verb_masks, dtype=torch.bool)
+            cl_masks = torch.tensor(cl_masks, dtype=torch.bool)  
+
+        else: 
+            sentences = text
+            images = image
+            targets = target
+            pad_masks = torch.zeros_like(text).masked_fill_(text == 0, 1).bool()
+
+        # Encoding images and text
+        vis = self.backbone.encode_image(images)
+        word, state = self.backbone.encode_text(sentences)
+
+        # FPN neck and decoder
+        fq, f5 = self.neck(vis, state)
+        b, c, h, w = fq.size()
+        fq = self.decoder(fq, word, pad_masks)
+        metric_tensor = fq  # b, c, h*w
+        fq = fq.reshape(b, c, h, w)
+
+        # Final prediction
+        pred = self.proj(fq, state)
+
+        if self.training:
+            if pred.shape[-2:] != targets.shape[-2:]:
+                targets = F.interpolate(targets, pred.shape[-2:], mode='nearest').detach()
+
+            loss = F.binary_cross_entropy_with_logits(pred[~cl_masks], targets[~cl_masks])
+
+            if self.metric_learning:
+                metric_loss = self.compute_metric_loss(metric_tensor, verb_masks, args=self.cfg)
+                loss = (loss + self.metric_loss_weight * metric_loss) / (1 + self.metric_loss_weight)
+            
+            return pred.detach(), targets, loss
+
+        return pred.detach()  # In eval mode, only return the predictions
+
+
+    def compute_metric_loss(self, metric_tensor, verb_mask, args) :
+        return None
+
+
+    def return_mask(self, emb_distance, positive_verbs, negative_verbs, verb_mask):
+        B_, B_ = emb_distance.shape
+        positive_mask = torch.zeros_like(emb_distance)
+        negative_mask = torch.ones_like(emb_distance)
+        positive_mask.fill_diagonal_(1)
+
+        if B_ < len(verb_mask):
+            # Considering only verbs that pass the verb_mask filter
+            positive_verbs = torch.tensor(positive_verbs)[verb_mask]
+            negative_verbs = torch.tensor(negative_verbs)[verb_mask]
+
+            # Exclude hard negatives from both masks (diagonal)
+            for i in range(B_):
+                if negative_verbs[i] == 1:
+                    positive_mask[i, i] = 0
+                    negative_mask[i, i] = 0
+
+            i = 0
+            while i < B_:
+                if positive_verbs[i] == 1:  
+                    if i + 1 < B_ and positive_verbs[i + 1] == 1:
+                        positive_mask[i, i + 1] = 1
+                        positive_mask[i + 1, i] = 1
+                    i += 2 
+                else:
+                    i += 1
+        else:
+            # Exclude hard negatives from both masks (diagonal)
+            for i in range(B_):
+                if negative_verbs[i] == 1:
+                    positive_mask[i, i] = 0
+                    negative_mask[i, i] = 0
+
+            # Apply the positive pairs logic similarly as above
+            i = 0
+            while i < B_:
+                if positive_verbs[i] == 1 and i + 1 < B_ and positive_verbs[i + 1] == 1:
+                    positive_mask[i, i + 1] = 1
+                    positive_mask[i + 1, i] = 1
+                    i += 2
+                else:
+                    i += 1
+
+        negative_mask = negative_mask - positive_mask
+
+        return positive_mask, negative_mask
+
+
+    def UniAngularContrastLoss(self, total_fq, positive_verbs, negative_verbs, m=0.5, tau=0.05, verbonly=True, args=None):
+        """
+        Angular Margin Contrastive Loss function with mask visualization.
+        """
+        verb_mask = positive_verbs + negative_verbs 
+        
+        if verbonly: 
+            emb = torch.mean(total_fq[verb_mask], dim=-1) 
+        else:
+            emb = torch.mean(total_fq, dim=-1) # (B, C)
+
+        B_ = emb.shape[0]
+        # emb = F.normalize(emb, p=2, dim=1) 
+        emb_i = emb.unsqueeze(1).repeat(1, B_, 1) # (B_, B_, C) 
+        emb_j = emb.unsqueeze(0).repeat(B_, 1, 1) # (B_, B_, C)
+        sim = nn.CosineSimilarity(dim=-1, eps=1e-6)
+        sim_matrix = sim(emb_i, emb_j).reshape(B_, B_)  # (B_, B_)
+        sim_matrix = torch.clamp(sim_matrix, min=-0.9999, max=0.9999)
+
+        # Apply angular margin for positive pairs using return_mask
+        positive_mask, negative_mask = self.return_mask(sim_matrix, positive_verbs, negative_verbs, verb_mask)
+
+        # Apply margin to positive pairs
+        sim_matrix_with_margin = sim_matrix.clone()
+        sim_matrix_with_margin[positive_mask.bool()] = torch.cos(torch.acos(sim_matrix[positive_mask.bool()]) + m / 57.2958)
+
+        # Scale logits with temperature
+        logits = sim_matrix_with_margin / tau
+
+        # Compute the softmax loss for all pairs
+        exp_logits = torch.exp(logits)
+        pos_exp_logits = exp_logits[positive_mask.bool()]
+        total_exp_logits = exp_logits.sum(dim=-1)
+
+        # Compute the final loss: L_arc = -log(e^(cos(theta + m)/tau) / sum(e^(cos(theta)/tau)))
+        positive_loss = -torch.log(pos_exp_logits / total_exp_logits[positive_mask.bool()])
+        angular_loss = positive_loss.mean()
+
+        return angular_loss