| 20 | |
| 21 | |
| 22 | class BiRefNet(nn.Module): |
| 23 | def __init__(self, bb_pretrained=True): |
| 24 | super(BiRefNet, self).__init__() |
| 25 | self.config = Config() |
| 26 | self.epoch = 1 |
| 27 | self.bb = build_backbone(self.config.bb, pretrained=bb_pretrained) |
| 28 | |
| 29 | channels = self.config.lateral_channels_in_collection |
| 30 | |
| 31 | if self.config.auxiliary_classification: |
| 32 | self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) |
| 33 | self.cls_head = nn.Sequential( |
| 34 | nn.Linear(channels[0], len(class_labels_TR_sorted)) |
| 35 | ) |
| 36 | |
| 37 | if self.config.squeeze_block: |
| 38 | self.squeeze_module = nn.Sequential(*[ |
| 39 | eval(self.config.squeeze_block.split('_x')[0])(channels[0]+sum(self.config.cxt), channels[0]) |
| 40 | for _ in range(eval(self.config.squeeze_block.split('_x')[1])) |
| 41 | ]) |
| 42 | |
| 43 | self.decoder = Decoder(channels) |
| 44 | |
| 45 | if self.config.ender: |
| 46 | self.dec_end = nn.Sequential( |
| 47 | nn.Conv2d(1, 16, 3, 1, 1), |
| 48 | nn.Conv2d(16, 1, 3, 1, 1), |
| 49 | nn.ReLU(inplace=True), |
| 50 | ) |
| 51 | |
| 52 | # refine patch-level segmentation |
| 53 | if self.config.refine: |
| 54 | if self.config.refine == 'itself': |
| 55 | self.stem_layer = StemLayer(in_channels=3+1, inter_channels=48, out_channels=3, norm_layer='BN' if self.config.batch_size > 1 else 'LN') |
| 56 | else: |
| 57 | self.refiner = eval('{}({})'.format(self.config.refine, 'in_channels=3+1')) |
| 58 | |
| 59 | if self.config.freeze_bb: |
| 60 | # Freeze the backbone... |
| 61 | print(self.named_parameters()) |
| 62 | for key, value in self.named_parameters(): |
| 63 | if 'bb.' in key and 'refiner.' not in key: |
| 64 | value.requires_grad = False |
| 65 | |
| 66 | def forward_enc(self, x): |
| 67 | if self.config.bb in ['vgg16', 'vgg16bn', 'resnet50']: |
| 68 | x1 = self.bb.conv1(x); x2 = self.bb.conv2(x1); x3 = self.bb.conv3(x2); x4 = self.bb.conv4(x3) |
| 69 | else: |
| 70 | x1, x2, x3, x4 = self.bb(x) |
| 71 | if self.config.mul_scl_ipt == 'cat': |
| 72 | B, C, H, W = x.shape |
| 73 | x1_, x2_, x3_, x4_ = self.bb(F.interpolate(x, size=(H//2, W//2), mode='bilinear', align_corners=True)) |
| 74 | x1 = torch.cat([x1, F.interpolate(x1_, size=x1.shape[2:], mode='bilinear', align_corners=True)], dim=1) |
| 75 | x2 = torch.cat([x2, F.interpolate(x2_, size=x2.shape[2:], mode='bilinear', align_corners=True)], dim=1) |
| 76 | x3 = torch.cat([x3, F.interpolate(x3_, size=x3.shape[2:], mode='bilinear', align_corners=True)], dim=1) |
| 77 | x4 = torch.cat([x4, F.interpolate(x4_, size=x4.shape[2:], mode='bilinear', align_corners=True)], dim=1) |
| 78 | elif self.config.mul_scl_ipt == 'add': |
| 79 | B, C, H, W = x.shape |
no outgoing calls
no test coverage detected