feat: add CondNet

README.md

@@ -83,6 +83,7 @@ Supported Heads/Methods:
 * [UPerNet][upernet]
 * [SFNet][sfnet] 
 * [SegFormer][segformer]
+* [CondNet][condnet]
 
 Supported Standalone Models:
 * [DDRNet][ddrnet]

models/heads/__init__.py

@@ -4,5 +4,6 @@
 from .fpn import FPNHead
 from .fapn import FaPNHead
 from .fcn import FCNHead
+from .condnet import CondHead
 
-__all__ = ['UPerHead', 'SegFormerHead', 'SFHead', 'FPNHead', 'FaPNHead', 'FCNHead']
+__all__ = ['UPerHead', 'SegFormerHead', 'SFHead', 'FPNHead', 'FaPNHead', 'FCNHead', 'CondHead']

models/heads/condnet.py

@@ -0,0 +1,67 @@
+import torch
+from torch import nn, Tensor
+from torch.nn import functional as F
+
+
+class ConvModule(nn.Sequential):
+    def __init__(self, c1, c2, k, s=1, p=0, d=1, g=1):
+        super().__init__(
+            nn.Conv2d(c1, c2, k, s, p, d, g, bias=False),
+            nn.BatchNorm2d(c2),
+            nn.ReLU(True)
+        )
+
+
+class CondHead(nn.Module):
+    def __init__(self, in_channel: int = 2048, channel: int = 512, num_classes: int = 19):
+        super().__init__()
+        self.num_classes = num_classes
+        self.weight_num = channel * num_classes
+        self.bias_num = num_classes
+
+        self.conv = ConvModule(in_channel, channel, 1)
+        self.dropout = nn.Dropout2d(0.1)
+
+        self.guidance_project = nn.Conv2d(channel, num_classes, 1)
+        self.filter_project = nn.Conv2d(channel*num_classes, self.weight_num + self.bias_num, 1, groups=num_classes)
+
+    def forward(self, features) -> Tensor:
+        x = self.dropout(self.conv(features[-1]))
+        B, C, H, W = x.shape
+        guidance_mask = self.guidance_project(x)
+        cond_logit = guidance_mask
+
+        key = x
+        value = x
+        guidance_mask = guidance_mask.softmax(dim=1).view(*guidance_mask.shape[:2], -1)
+        key = key.view(B, C, -1).permute(0, 2, 1)
+
+        cond_filters = torch.matmul(guidance_mask, key)
+        cond_filters /= H * W
+        cond_filters = cond_filters.view(B, -1, 1, 1)
+        cond_filters = self.filter_project(cond_filters)
+        cond_filters = cond_filters.view(B, -1)
+
+        weight, bias = torch.split(cond_filters, [self.weight_num, self.bias_num], dim=1)
+        weight = weight.reshape(B * self.num_classes, -1, 1, 1)
+        bias = bias.reshape(B * self.num_classes)
+
+        value = value.view(-1, H, W).unsqueeze(0)
+        seg_logit = F.conv2d(value, weight, bias, 1, 0, groups=B).view(B, self.num_classes, H, W)
+
+        if self.training:
+            return cond_logit, seg_logit
+        return seg_logit
+
+
+if __name__ == '__main__':
+    import sys
+    sys.path.insert(0, '.')
+    from models.backbones.resnetd import ResNetD
+    backbone = ResNetD('50')
+    head = CondHead()
+    x = torch.randn(2, 3, 224, 224)
+    features = backbone(x)
+    outs = head(features)
+    for out in outs:
+        print(out.shape)