Explorar o código

Allow mask not in sample in ReloadMask

Bobholamovic %!s(int64=2) %!d(string=hai) anos
pai
achega
739888e6ca
Modificáronse 1 ficheiros con 3 adicións e 2 borrados
  1. 3 2
      paddlers/transforms/operators.py

+ 3 - 2
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