Remove redundant part of our_ref in inference.

app.py

@@ -35,7 +35,7 @@ class ImagePreprocessor():
         return image
 
 
-model = BiRefNet().to(device)
+model = BiRefNet(bb_pretrained=False).to(device)
 state_dict = './BiRefNet_ep580.pth'
 if os.path.exists(state_dict):
     birefnet_dict = torch.load(state_dict, map_location=device)

models/baseline.py

@@ -20,11 +20,11 @@ from models.refinement.stem_layer import StemLayer
 
 
 class BiRefNet(nn.Module):
-    def __init__(self):
+    def __init__(self, bb_pretrained=True):
         super(BiRefNet, self).__init__()
         self.config = Config()
         self.epoch = 1
-        self.bb = build_backbone(self.config.bb, pretrained=False)
+        self.bb = build_backbone(self.config.bb, pretrained=bb_pretrained)
 
         channels = self.config.lateral_channels_in_collection
 
@@ -126,7 +126,7 @@ class BiRefNet(nn.Module):
             x4 = self.squeeze_module(x4)
         ########## Decoder ##########
         features = [x, x1, x2, x3, x4]
-        if self.config.out_ref:
+        if self.training and self.config.out_ref:
             features.append(laplacian(torch.mean(x, dim=1).unsqueeze(1), kernel_size=5))
         scaled_preds = self.decoder(features)
         return scaled_preds, class_preds
@@ -231,7 +231,7 @@ class Decoder(nn.Module):
         return torch.cat(patches_batch, dim=0)
 
     def forward(self, features):
-        if self.config.out_ref:
+        if self.training and self.config.out_ref:
             outs_gdt_pred = []
             outs_gdt_label = []
             x, x1, x2, x3, x4, gdt_gt = features
@@ -249,18 +249,19 @@ class Decoder(nn.Module):
         p3 = self.decoder_block3(_p3)
         m3 = self.conv_ms_spvn_3(p3) if self.config.ms_supervision else None
         if self.config.out_ref:
-            # >> GT:
-            # m3 --dilation--> m3_dia
-            # G_3^gt * m3_dia --> G_3^m, which is the label of gradient
-            m3_dia = m3
-            gdt_label_main_3 = gdt_gt * F.interpolate(m3_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True)
-            outs_gdt_label.append(gdt_label_main_3)
-            # >> Pred:
-            # p3 --conv--BN--> F_3^G, where F_3^G predicts the \hat{G_3} with xx
-            # F_3^G --sigmoid--> A_3^G
             p3_gdt = self.gdt_convs_3(p3)
-            gdt_pred_3 = self.gdt_convs_pred_3(p3_gdt)
-            outs_gdt_pred.append(gdt_pred_3)
+            if self.training:
+                # >> GT:
+                # m3 --dilation--> m3_dia
+                # G_3^gt * m3_dia --> G_3^m, which is the label of gradient
+                m3_dia = m3
+                gdt_label_main_3 = gdt_gt * F.interpolate(m3_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True)
+                outs_gdt_label.append(gdt_label_main_3)
+                # >> Pred:
+                # p3 --conv--BN--> F_3^G, where F_3^G predicts the \hat{G_3} with xx
+                # F_3^G --sigmoid--> A_3^G
+                gdt_pred_3 = self.gdt_convs_pred_3(p3_gdt)
+                outs_gdt_pred.append(gdt_pred_3)
             gdt_attn_3 = self.gdt_convs_attn_3(p3_gdt).sigmoid()
             # >> Finally:
             # p3 = p3 * A_3^G
@@ -274,14 +275,15 @@ class Decoder(nn.Module):
         p2 = self.decoder_block2(_p2)
         m2 = self.conv_ms_spvn_2(p2) if self.config.ms_supervision else None
         if self.config.out_ref:
-            # >> GT:
-            m2_dia = m2
-            gdt_label_main_2 = gdt_gt * F.interpolate(m2_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True)
-            outs_gdt_label.append(gdt_label_main_2)
-            # >> Pred:
             p2_gdt = self.gdt_convs_2(p2)
-            gdt_pred_2 = self.gdt_convs_pred_2(p2_gdt)
-            outs_gdt_pred.append(gdt_pred_2)
+            if self.training:
+                # >> GT:
+                m2_dia = m2
+                gdt_label_main_2 = gdt_gt * F.interpolate(m2_dia, size=gdt_gt.shape[2:], mode='bilinear', align_corners=True)
+                outs_gdt_label.append(gdt_label_main_2)
+                # >> Pred:
+                gdt_pred_2 = self.gdt_convs_pred_2(p2_gdt)
+                outs_gdt_pred.append(gdt_pred_2)
             gdt_attn_2 = self.gdt_convs_attn_2(p2_gdt).sigmoid()
             # >> Finally:
             p2 = p2 * gdt_attn_2