[Fix] Fix binary C=1 focal loss & dataset fileio (open-mmlab#2935)

nahidnazifi87 · Apr 23, 2023 · da09304 · da09304
1 parent 8327e29
commit da09304
Show file tree

Hide file tree

Showing 5 changed files with 25 additions and 6 deletions.
diff --git a/mmseg/datasets/chase_db1.py b/mmseg/datasets/chase_db1.py
@@ -1,4 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import mmengine.fileio as fileio
 
 from mmseg.registry import DATASETS
 from .basesegdataset import BaseSegDataset
@@ -27,4 +28,5 @@ def __init__(self,
             seg_map_suffix=seg_map_suffix,
             reduce_zero_label=reduce_zero_label,
             **kwargs)
-        assert self.file_client.exists(self.data_prefix['img_path'])
+        assert fileio.exists(
+            self.data_prefix['img_path'], backend_args=self.backend_args)
diff --git a/mmseg/datasets/drive.py b/mmseg/datasets/drive.py
@@ -1,4 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import mmengine.fileio as fileio
 
 from mmseg.registry import DATASETS
 from .basesegdataset import BaseSegDataset
@@ -27,4 +28,5 @@ def __init__(self,
             seg_map_suffix=seg_map_suffix,
             reduce_zero_label=reduce_zero_label,
             **kwargs)
-        assert self.file_client.exists(self.data_prefix['img_path'])
+        assert fileio.exists(
+            self.data_prefix['img_path'], backend_args=self.backend_args)
diff --git a/mmseg/datasets/hrf.py b/mmseg/datasets/hrf.py
@@ -1,4 +1,5 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import mmengine.fileio as fileio
 
 from mmseg.registry import DATASETS
 from .basesegdataset import BaseSegDataset
@@ -27,4 +28,5 @@ def __init__(self,
             seg_map_suffix=seg_map_suffix,
             reduce_zero_label=reduce_zero_label,
             **kwargs)
-        assert self.file_client.exists(self.data_prefix['img_path'])
+        assert fileio.exists(
+            self.data_prefix['img_path'], backend_args=self.backend_args)
diff --git a/mmseg/datasets/stare.py b/mmseg/datasets/stare.py
@@ -1,4 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import mmengine.fileio as fileio
+
 from mmseg.registry import DATASETS
 from .basesegdataset import BaseSegDataset
 
@@ -26,4 +28,5 @@ def __init__(self,
             seg_map_suffix=seg_map_suffix,
             reduce_zero_label=reduce_zero_label,
             **kwargs)
-        assert self.file_client.exists(self.data_prefix['img_path'])
+        assert fileio.exists(
+            self.data_prefix['img_path'], backend_args=self.backend_args)
diff --git a/mmseg/models/losses/focal_loss.py b/mmseg/models/losses/focal_loss.py
@@ -271,7 +271,13 @@ def forward(self,
             num_classes = pred.size(1)
             if torch.cuda.is_available() and pred.is_cuda:
                 if target.dim() == 1:
-                    one_hot_target = F.one_hot(target, num_classes=num_classes)
+                    one_hot_target = F.one_hot(
+                        target, num_classes=num_classes + 1)
+                    if num_classes == 1:
+                        one_hot_target = one_hot_target[:, 1]
+                        target = 1 - target
+                    else:
+                        one_hot_target = one_hot_target[:, :num_classes]
                 else:
                     one_hot_target = target
                     target = target.argmax(dim=1)
@@ -280,7 +286,11 @@ def forward(self,
             else:
                 one_hot_target = None
                 if target.dim() == 1:
-                    target = F.one_hot(target, num_classes=num_classes)
+                    target = F.one_hot(target, num_classes=num_classes + 1)
+                    if num_classes == 1:
+                        target = target[:, 1]
+                    else:
+                        target = target[:, num_classes]
                 else:
                     valid_mask = (target.argmax(dim=1) != ignore_index).view(
                         -1, 1)