fix: apply correct weights for static class balancing

Author: torms3

deepem/loss/affinity.py

@@ -56,19 +56,6 @@ def forward(self, preds, targets, masks):
 
         return loss, nmsk
 
-    # def class_balancing(self, target, mask):
-    #     if not self.balancing:
-    #         return mask
-    #     dtype = mask.type()
-    #     m_int = mask * torch.eq(target, 1).type(dtype)
-    #     m_ext = mask * torch.eq(target, 0).type(dtype)
-    #     n_int = m_int.sum().item()
-    #     n_ext = m_ext.sum().item()
-    #     if n_int > 0 and n_ext > 0:
-    #         m_int *= n_ext/(n_int + n_ext)
-    #         m_ext *= n_int/(n_int + n_ext)
-    #     return (m_int + m_ext).type(dtype)
-
 
 class AffinityLoss(nn.Module):
     def __init__(self, edges, criterion, size_average=False,

deepem/train/utils.py

@@ -20,6 +20,8 @@ def get_criteria(opt):
         weight1=opt.class_weight1,
     ) if opt.class_balancing else None
 
+    is_dynamic = (opt.class_weight0 is None) and (opt.class_weight1 is None)
+
     for k in opt.out_spec:
         if k == 'affinity' or k == 'long_range':
             if k == 'affinity':
@@ -42,7 +44,14 @@ def get_criteria(opt):
                 params['margin0'] = 0
                 params['margin1'] = 0
                 params['inverse'] = False
-            params['class_balancer'] = balancer
+            params['class_balancer'] = (
+                balancer
+                if is_dynamic
+                else BinaryWeightBalancer(
+                    weight0=opt.class_weight1,
+                    weight1=opt.class_weight0,
+                )
+            )
             criteria[k] = getattr(loss, 'BCELoss')(**params)
     return criteria