Allow mask not in sample in ReloadMask

2 years ago · 739888e6ca
parent 66894e3adb
commit 739888e6ca
1 changed files with 3 additions and 2 deletions
--- a/paddlers/transforms/operators.py
+++ b/paddlers/transforms/operators.py
@ -1943,8 +1943,9 @@ class RandomSwap(Transform):

 class ReloadMask(Transform):
    def apply(self, sample):
-        sample['mask'] = F.decode_seg_mask(sample['mask_ori'])
-        if 'aux_masks' in sample:
+        if 'mask' in sample or 'mask_ori' in sample:
+            sample['mask'] = F.decode_seg_mask(sample['mask_ori'])
+        if 'aux_masks' in sample or 'aux_masks_ori' in sample:
            sample['aux_masks'] = list(
                map(F.decode_seg_mask, sample['aux_masks_ori']))
        return sample