move get_class_weight to utils

mmseg/models/losses/cross_entropy_loss.py

@@ -3,7 +3,7 @@
 import torch.nn.functional as F
 
 from ..builder import LOSSES
-from .utils import weight_reduce_loss
+from .utils import get_class_weight, weight_reduce_loss
 
 
 def cross_entropy(pred,
@@ -163,7 +163,7 @@ def __init__(self,
         self.use_mask = use_mask
         self.reduction = reduction
         self.loss_weight = loss_weight
-        self.class_weight = self._get_class_weight(class_weight)
+        self.class_weight = get_class_weight(class_weight)
 
         if self.use_sigmoid:
             self.cls_criterion = binary_cross_entropy
@@ -172,28 +172,6 @@ def __init__(self,
         else:
             self.cls_criterion = cross_entropy
 
-    def _get_class_weight(self, class_weight):
-        """Get class weight for loss function.
-
-        Args:
-            class_weight (list[float] | str | None): If class_weight is a str,
-                take it as a file name and read from it.
-        """
-        import numpy as np
-        import mmcv
-
-        if isinstance(class_weight, str):
-            # take it as a file path
-            if class_weight.endswith('.npy'):
-                class_weight = np.load(class_weight)
-            else:
-                file_format = class_weight.split('.')[-1]
-                assert file_format in ['json', 'yaml', 'pkl'], \
-                    f'unsupported class_weight file type {file_format}'
-                class_weight = mmcv.load(class_weight)
-
-        return class_weight
-
     def forward(self,
                 cls_score,
                 label,

mmseg/models/losses/dice_loss.py

@@ -5,7 +5,7 @@
 import torch.nn.functional as F
 
 from ..builder import LOSSES
-from .utils import weighted_loss
+from .utils import get_class_weight, weighted_loss
 
 
 @weighted_loss
@@ -81,32 +81,10 @@ def __init__(self,
         self.smooth = smooth
         self.exponent = exponent
         self.reduction = reduction
-        self.class_weight = self._get_class_weight(class_weight)
+        self.class_weight = get_class_weight(class_weight)
         self.loss_weight = loss_weight
         self.ignore_index = ignore_index
 
-    def _get_class_weight(self, class_weight):
-        """Get class weight for loss function.
-
-        Args:
-            class_weight (list[float] | str | None): If class_weight is a str,
-                take it as a file name and read from it.
-        """
-        import numpy as np
-        import mmcv
-
-        if isinstance(class_weight, str):
-            # take it as a file path
-            if class_weight.endswith('.npy'):
-                class_weight = np.load(class_weight)
-            else:
-                file_format = class_weight.split('.')[-1]
-                assert file_format in ['json', 'yaml', 'pkl'], \
-                    f'unsupported class_weight file type {file_format}'
-                class_weight = mmcv.load(class_weight)
-
-        return class_weight
-
     def forward(self,
                 pred,
                 target,

mmseg/models/losses/lovasz_loss.py

@@ -8,7 +8,7 @@
 import torch.nn.functional as F
 
 from ..builder import LOSSES
-from .utils import weight_reduce_loss
+from .utils import get_class_weight, weight_reduce_loss
 
 
 def lovasz_grad(gt_sorted):
@@ -269,29 +269,7 @@ def __init__(self,
         self.per_image = per_image
         self.reduction = reduction
         self.loss_weight = loss_weight
-        self.class_weight = self._get_class_weight(class_weight)
-
-    def _get_class_weight(self, class_weight):
-        """Get class weight for loss function.
-
-        Args:
-            class_weight (list[float] | str | None): If class_weight is a str,
-                take it as a file name and read from it.
-        """
-        import numpy as np
-        import mmcv
-
-        if isinstance(class_weight, str):
-            # take it as a file path
-            if class_weight.endswith('.npy'):
-                class_weight = np.load(class_weight)
-            else:
-                file_format = class_weight.split('.')[-1]
-                assert file_format in ['json', 'yaml', 'pkl'], \
-                    f'unsupported class_weight file type {file_format}'
-                class_weight = mmcv.load(class_weight)
-
-        return class_weight
+        self.class_weight = get_class_weight(class_weight)
 
     def forward(self,
                 cls_score,

mmseg/models/losses/utils.py

@@ -1,8 +1,28 @@
 import functools
 
+import mmcv
+import numpy as np
 import torch.nn.functional as F
 
 
+def get_class_weight(class_weight):
+    """Get class weight for loss function.
+
+    Args:
+        class_weight (list[float] | str | None): If class_weight is a str,
+            take it as a file name and read from it.
+    """
+    if isinstance(class_weight, str):
+        # take it as a file path
+        if class_weight.endswith('.npy'):
+            class_weight = np.load(class_weight)
+        else:
+            # pkl, json or yaml
+            class_weight = mmcv.load(class_weight)
+
+    return class_weight
+
+
 def reduce_loss(loss, reduction):
     """Reduce loss as specified.