Move all BiRefNet github codes to the first level directory.

__init__.py

@@ -0,0 +1,6 @@
+from os.path import dirname, basename, isfile, join
+import glob
+
+
+modules = glob.glob(join(dirname(__file__), "*.py"))
+__all__ = [ basename(f)[:-3] for f in modules if isfile(f) and not f.endswith('__init__.py')]

models/backbones/__init__.py

@@ -0,0 +1,6 @@
+from os.path import dirname, basename, isfile, join
+import glob
+
+
+modules = glob.glob(join(dirname(__file__), "*.py"))
+__all__ = [ basename(f)[:-3] for f in modules if isfile(f) and not f.endswith('__init__.py')]

models/modules/__init__.py

@@ -0,0 +1,6 @@
+from os.path import dirname, basename, isfile, join
+import glob
+
+
+modules = glob.glob(join(dirname(__file__), "*.py"))
+__all__ = [ basename(f)[:-3] for f in modules if isfile(f) and not f.endswith('__init__.py')]

models/modules/refinement/refiner.py

@@ -0,0 +1,253 @@
+import torch
+import torch.nn as nn
+from collections import OrderedDict
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision.models import vgg16, vgg16_bn
+from torchvision.models import resnet50
+
+from config import Config
+from dataset import class_labels_TR_sorted
+from models.backbones.build_backbone import build_backbone
+from models.modules.decoder_blocks import BasicDecBlk
+from models.modules.lateral_blocks import BasicLatBlk
+from models.modules.ing import *
+from models.refinement.stem_layer import StemLayer
+
+
+class RefinerPVTInChannels4(nn.Module):
+    def __init__(self, in_channels=3+1):
+        super(RefinerPVTInChannels4, self).__init__()
+        self.config = Config()
+        self.epoch = 1
+        self.bb = build_backbone(self.config.bb, params_settings='in_channels=4')
+
+        lateral_channels_in_collection = {
+            'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64],
+            'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64],
+            'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192],
+        }
+        channels = lateral_channels_in_collection[self.config.bb]
+        self.squeeze_module = BasicDecBlk(channels[0], channels[0])
+
+        self.decoder = Decoder(channels)
+
+        if 0:
+            for key, value in self.named_parameters():
+                if 'bb.' in key:
+                    value.requires_grad = False
+
+    def forward(self, x):
+        if isinstance(x, list):
+            x = torch.cat(x, dim=1)
+        ########## Encoder ##########
+        if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']:
+            x1 = self.bb.conv1(x)
+            x2 = self.bb.conv2(x1)
+            x3 = self.bb.conv3(x2)
+            x4 = self.bb.conv4(x3)
+        else:
+            x1, x2, x3, x4 = self.bb(x)
+
+        x4 = self.squeeze_module(x4)
+
+        ########## Decoder ##########
+
+        features = [x, x1, x2, x3, x4]
+        scaled_preds = self.decoder(features)
+
+        return scaled_preds
+
+
+class Refiner(nn.Module):
+    def __init__(self, in_channels=3+1):
+        super(Refiner, self).__init__()
+        self.config = Config()
+        self.epoch = 1
+        self.stem_layer = StemLayer(in_channels=in_channels, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN')
+        self.bb = build_backbone(self.config.bb)
+
+        lateral_channels_in_collection = {
+            'vgg16': [512, 256, 128, 64], 'vgg16bn': [512, 256, 128, 64], 'resnet50': [1024, 512, 256, 64],
+            'pvt_v2_b2': [512, 320, 128, 64], 'pvt_v2_b5': [512, 320, 128, 64],
+            'swin_v1_b': [1024, 512, 256, 128], 'swin_v1_l': [1536, 768, 384, 192],
+        }
+        channels = lateral_channels_in_collection[self.config.bb]
+        self.squeeze_module = BasicDecBlk(channels[0], channels[0])
+
+        self.decoder = Decoder(channels)
+
+        if 0:
+            for key, value in self.named_parameters():
+                if 'bb.' in key:
+                    value.requires_grad = False
+
+    def forward(self, x):
+        if isinstance(x, list):
+            x = torch.cat(x, dim=1)
+        x = self.stem_layer(x)
+        ########## Encoder ##########
+        if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']:
+            x1 = self.bb.conv1(x)
+            x2 = self.bb.conv2(x1)
+            x3 = self.bb.conv3(x2)
+            x4 = self.bb.conv4(x3)
+        else:
+            x1, x2, x3, x4 = self.bb(x)
+
+        x4 = self.squeeze_module(x4)
+
+        ########## Decoder ##########
+
+        features = [x, x1, x2, x3, x4]
+        scaled_preds = self.decoder(features)
+
+        return scaled_preds
+
+
+class Decoder(nn.Module):
+    def __init__(self, channels):
+        super(Decoder, self).__init__()
+        self.config = Config()
+        DecoderBlock = eval('BasicDecBlk')
+        LateralBlock = eval('BasicLatBlk')
+
+        self.decoder_block4 = DecoderBlock(channels[0], channels[1])
+        self.decoder_block3 = DecoderBlock(channels[1], channels[2])
+        self.decoder_block2 = DecoderBlock(channels[2], channels[3])
+        self.decoder_block1 = DecoderBlock(channels[3], channels[3]//2)
+
+        self.lateral_block4 = LateralBlock(channels[1], channels[1])
+        self.lateral_block3 = LateralBlock(channels[2], channels[2])
+        self.lateral_block2 = LateralBlock(channels[3], channels[3])
+
+        if self.config.ms_supervision:
+            self.conv_ms_spvn_4 = nn.Conv2d(channels[1], 1, 1, 1, 0)
+            self.conv_ms_spvn_3 = nn.Conv2d(channels[2], 1, 1, 1, 0)
+            self.conv_ms_spvn_2 = nn.Conv2d(channels[3], 1, 1, 1, 0)
+        self.conv_out1 = nn.Sequential(nn.Conv2d(channels[3]//2, 1, 1, 1, 0))
+
+    def forward(self, features):
+        x, x1, x2, x3, x4 = features
+        outs = []
+        p4 = self.decoder_block4(x4)
+        _p4 = F.interpolate(p4, size=x3.shape[2:], mode='bilinear', align_corners=True)
+        _p3 = _p4 + self.lateral_block4(x3)
+
+        p3 = self.decoder_block3(_p3)
+        _p3 = F.interpolate(p3, size=x2.shape[2:], mode='bilinear', align_corners=True)
+        _p2 = _p3 + self.lateral_block3(x2)
+
+        p2 = self.decoder_block2(_p2)
+        _p2 = F.interpolate(p2, size=x1.shape[2:], mode='bilinear', align_corners=True)
+        _p1 = _p2 + self.lateral_block2(x1)
+
+        _p1 = self.decoder_block1(_p1)
+        _p1 = F.interpolate(_p1, size=x.shape[2:], mode='bilinear', align_corners=True)
+        p1_out = self.conv_out1(_p1)
+
+        if self.config.ms_supervision:
+            outs.append(self.conv_ms_spvn_4(p4))
+            outs.append(self.conv_ms_spvn_3(p3))
+            outs.append(self.conv_ms_spvn_2(p2))
+        outs.append(p1_out)
+        return outs
+
+
+class RefUNet(nn.Module):
+    # Refinement
+    def __init__(self, in_channels=3+1):
+        super(RefUNet, self).__init__()
+        self.encoder_1 = nn.Sequential(
+            nn.Conv2d(in_channels, 64, 3, 1, 1),
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True)
+        )
+
+        self.encoder_2 = nn.Sequential(
+            nn.MaxPool2d(2, 2, ceil_mode=True),
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True)
+        )
+
+        self.encoder_3 = nn.Sequential(
+            nn.MaxPool2d(2, 2, ceil_mode=True),
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True)
+        )
+
+        self.encoder_4 = nn.Sequential(
+            nn.MaxPool2d(2, 2, ceil_mode=True),
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True)
+        )
+
+        self.pool4 = nn.MaxPool2d(2, 2, ceil_mode=True)
+        #####
+        self.decoder_5 = nn.Sequential(
+            nn.Conv2d(64, 64, 3, 1, 1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True)
+        )
+        #####
+        self.decoder_4 = nn.Sequential(
+            nn.Conv2d(128, 64, 3, 1, 1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True)
+        )
+
+        self.decoder_3 = nn.Sequential(
+            nn.Conv2d(128, 64, 3, 1, 1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True)
+        )
+
+        self.decoder_2 = nn.Sequential(
+            nn.Conv2d(128, 64, 3, 1, 1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True)
+        )
+
+        self.decoder_1 = nn.Sequential(
+            nn.Conv2d(128, 64, 3, 1, 1),
+            nn.BatchNorm2d(64),
+            nn.ReLU(inplace=True)
+        )
+
+        self.conv_d0 = nn.Conv2d(64, 1, 3, 1, 1)
+
+        self.upscore2 = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+
+    def forward(self, x):
+        outs = []
+        if isinstance(x, list):
+            x = torch.cat(x, dim=1)
+        hx = x
+
+        hx1 = self.encoder_1(hx)
+        hx2 = self.encoder_2(hx1)
+        hx3 = self.encoder_3(hx2)
+        hx4 = self.encoder_4(hx3)
+
+        hx = self.decoder_5(self.pool4(hx4))
+        hx = torch.cat((self.upscore2(hx), hx4), 1)
+
+        d4 = self.decoder_4(hx)
+        hx = torch.cat((self.upscore2(d4), hx3), 1)
+
+        d3 = self.decoder_3(hx)
+        hx = torch.cat((self.upscore2(d3), hx2), 1)
+
+        d2 = self.decoder_2(hx)
+        hx = torch.cat((self.upscore2(d2), hx1), 1)
+
+        d1 = self.decoder_1(hx)
+
+        x = self.conv_d0(d1)
+        outs.append(x)
+        return outs

models/modules/refinement/stem_layer.py

@@ -0,0 +1,45 @@
+import torch.nn as nn
+from models.modules.utils import build_act_layer, build_norm_layer
+
+
+class StemLayer(nn.Module):
+    r""" Stem layer of InternImage
+    Args:
+        in_channels (int): number of input channels
+        out_channels (int): number of output channels
+        act_layer (str): activation layer
+        norm_layer (str): normalization layer
+    """
+
+    def __init__(self,
+                 in_channels=3+1,
+                 inter_channels=48,
+                 out_channels=96,
+                 act_layer='GELU',
+                 norm_layer='BN'):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_channels,
+                               inter_channels,
+                               kernel_size=3,
+                               stride=1,
+                               padding=1)
+        self.norm1 = build_norm_layer(
+            inter_channels, norm_layer, 'channels_first', 'channels_first'
+        )
+        self.act = build_act_layer(act_layer)
+        self.conv2 = nn.Conv2d(inter_channels,
+                               out_channels,
+                               kernel_size=3,
+                               stride=1,
+                               padding=1)
+        self.norm2 = build_norm_layer(
+            out_channels, norm_layer, 'channels_first', 'channels_first'
+        )
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.norm1(x)
+        x = self.act(x)
+        x = self.conv2(x)
+        x = self.norm2(x)
+        return x

models/refinement/__init__.py

@@ -0,0 +1,6 @@
+from os.path import dirname, basename, isfile, join
+import glob
+
+
+modules = glob.glob(join(dirname(__file__), "*.py"))
+__all__ = [ basename(f)[:-3] for f in modules if isfile(f) and not f.endswith('__init__.py')]