Add option of structure loss and link to single image inference on colab to README.

Author: ZhengPeng7

README.md

@@ -49,6 +49,8 @@ Choose the one you like to try with clicks instead of codes:
    + Thanks [**viperyl/ComfyUI-BiRefNet**](https://github.com/viperyl/ComfyUI-BiRefNet): this project packs BiRefNet as **ComfyUI nodes**, and makes this SOTA model easier use for everyone.
      <p align="center"><img src="https://drive.google.com/thumbnail?id=1KfxCQUUa2y9T-aysEaeVVjCUt3Z0zSkL&sz=w1620" /></p>
 
+   + Thanks [**Rishabh**](https://github.com/rishabh063) for offerring a demo for the [easier single image inference on colab](https://colab.research.google.com/drive/14Dqg7oeBkFEtchaHLNpig2BcdkZEogba?usp=drive_link).
+
 2. **More Visual Comparisons**
    + Thanks [**twitter.com/ZHOZHO672070**](https://twitter.com/ZHOZHO672070) for the comparison with more background-removal methods in images:
 

config.py

@@ -13,7 +13,7 @@ def __init__(self) -> None:
         self.training_set = {
             'DIS5K': 'DIS-TR',
             'COD': 'TR-COD10K+TR-CAMO',
-            'HRSOD': ['TR-DUTS', 'TR-HRSOD+TR-UHRSD', 'TR-DUTS+TR-HRSOD+TR-UHRSD'][1],
+            'HRSOD': ['TR-DUTS', 'TR-DUTS+TR-HRSOD', 'TR-DUTS+TR-UHRSD', 'TR-HRSOD+TR-UHRSD', 'TR-DUTS+TR-HRSOD+TR-UHRSD'][3],
             'DIS5K+HRSOD+HRS10K': 'DIS-TE1+DIS-TE2+DIS-TE3+DIS-TE4+DIS-TR+TE-HRS10K+TE-HRSOD+TE-UHRSD+TR-HRS10K+TR-HRSOD+TR-UHRSD',     # leave DIS-VD for evaluation.
             'P3M-10k': 'TR-P3M-10k',
         }[self.task]
@@ -102,6 +102,7 @@ def __init__(self) -> None:
             'reg': 100 * 0,
             'ssim': 10 * 1,          # help contours,
             'cnt': 5 * 0,          # help contours
+            'structure': 5 * 0,    # structure loss
         }
         self.lambdas_cls = {
             'ce': 5.0

eval_existingOnes.py

@@ -121,7 +121,7 @@ def do_eval(opt):
     opt = parser.parse_args()
 
     os.makedirs(opt.save_dir, exist_ok=True)
-    opt.model_lst = [m for m in sorted(os.listdir(opt.pred_root), key=lambda x: int(x.split('ep')[-1]), reverse=True) if int(m.split('ep')[-1]) % 1 == 0]
+    opt.model_lst = [m for m in sorted(os.listdir(opt.pred_root), key=lambda x: int(x.split('epoch_')[-1]), reverse=True) if int(m.split('epoch_')[-1]) % 1 == 0]
 
     # check the integrity of each candidates
     if opt.check_integrity:

inference.py

@@ -54,7 +54,7 @@ def main(args):
         model = BiRefNet(bb_pretrained=False)
     weights_lst = sorted(
         glob(os.path.join(args.ckpt_folder, '*.pth')) if args.ckpt_folder else [args.ckpt],
-        key=lambda x: int(x.split('ep')[-1].split('.pth')[0]),
+        key=lambda x: int(x.split('epoch_')[-1].split('.pth')[0]),
         reverse=True
     )
     for testset in args.testsets.split('+'):
@@ -64,7 +64,7 @@ def main(args):
             batch_size=config.batch_size_valid, shuffle=False, num_workers=config.num_workers, pin_memory=True
         )
         for weights in weights_lst:
-            if int(weights.strip('.pth').split('ep')[-1]) % 1 != 0:
+            if int(weights.strip('.pth').split('epoch_')[-1]) % 1 != 0:
                 continue
             print('\tInferencing {}...'.format(weights))
             # model.load_state_dict(torch.load(weights, map_location='cpu'))

loss.py

@@ -86,6 +86,23 @@ def forward(self, pred, target):
         return IoU
 
 
+class StructureLoss(torch.nn.Module):
+    def __init__(self):
+        super(StructureLoss, self).__init__()
+
+    def forward(self, pred, target):
+        weit  = 1+5*torch.abs(F.avg_pool2d(target, kernel_size=31, stride=1, padding=15)-target)
+        wbce  = F.binary_cross_entropy_with_logits(pred, target, reduction='none')
+        wbce  = (weit*wbce).sum(dim=(2,3))/weit.sum(dim=(2,3))
+
+        pred  = torch.sigmoid(pred)
+        inter = ((pred * target) * weit).sum(dim=(2, 3))
+        union = ((pred + target) * weit).sum(dim=(2, 3))
+        wiou  = 1-(inter+1)/(union-inter+1)
+
+        return (wbce+wiou).mean()
+
+
 class PatchIoULoss(torch.nn.Module):
     def __init__(self):
         super(PatchIoULoss, self).__init__()
@@ -158,15 +175,16 @@ def __init__(self):
             self.criterions_last['reg'] = ThrReg_loss()
         if 'cnt' in self.lambdas_pix_last and self.lambdas_pix_last['cnt']:
             self.criterions_last['cnt'] = ContourLoss()
+        if 'structure' in self.lambdas_pix_last and self.lambdas_pix_last['structure']:
+            self.criterions_last['structure'] = StructureLoss()
 
     def forward(self, scaled_preds, gt):
         loss = 0.
         for _, pred_lvl in enumerate(scaled_preds):
             if pred_lvl.shape != gt.shape:
                 pred_lvl = nn.functional.interpolate(pred_lvl, size=gt.shape[2:], mode='bilinear', align_corners=True)
-            pred_lvl = pred_lvl.sigmoid()
             for criterion_name, criterion in self.criterions_last.items():
-                _loss = criterion(pred_lvl, gt) * self.lambdas_pix_last[criterion_name]
+                _loss = criterion(pred_lvl.sigmoid() if criterion_name not in ('structure',) else pred_lvl, gt) * self.lambdas_pix_last[criterion_name]
                 loss += _loss
                 # print(criterion_name, _loss.item())
         return loss

train.py

@@ -101,7 +101,7 @@ def init_models_optimizers(epochs, to_be_distributed):
             state_dict = torch.load(args.resume, map_location='cpu')
             state_dict = check_state_dict(state_dict)
             model.load_state_dict(state_dict)
-            epoch_st = int(args.resume.rstrip('.pth').split('ep')[-1]) + 1
+            epoch_st = int(args.resume.rstrip('.pth').split('epoch_')[-1]) + 1
         else:
             logger.info("=> no checkpoint found at '{}'".format(args.resume))
     if to_be_distributed:
@@ -301,7 +301,7 @@ def main():
         if epoch >= args.epochs - config.save_last and epoch % config.save_step == 0:
             torch.save(
                 trainer.model.module.state_dict() if to_be_distributed else trainer.model.state_dict(),
-                os.path.join(args.ckpt_dir, 'ep{}.pth'.format(epoch))
+                os.path.join(args.ckpt_dir, 'epoch_{}.pth'.format(epoch))
             )
         if config.val_step and epoch >= args.epochs - config.save_last and (args.epochs - epoch) % config.val_step == 0:
             if to_be_distributed:

waiting4eval.py

@@ -111,7 +111,7 @@ def main():
         models_detected = [
             m for idx_m, m in enumerate(sorted(
                 glob(os.path.join(ckpt_dir, '*.pth')),
-                key=lambda x: int(x.rstrip('.pth').split('ep')[-1]), reverse=True
+                key=lambda x: int(x.rstrip('.pth').split('epoch_')[-1]), reverse=True
             )) if idx_m % args_eval.val_step == args_eval.program_id and m not in models_evaluated + models_evaluated_global
         ]
         if models_detected:
@@ -127,7 +127,7 @@ def main():
             # evaluate the current model
             state_dict = torch.load(model_not_evaluated_latest, map_location=device)
             model.load_state_dict(state_dict, strict=False)
-            validate_model(model, test_loaders, int(model_not_evaluated_latest.rstrip('.pth').split('ep')[-1]))
+            validate_model(model, test_loaders, int(model_not_evaluated_latest.rstrip('.pth').split('epoch_')[-1]))
             continous_sleep_time = 0
             print('Duration of this evaluation:', time() - time_st)
         else: