From 739888e6ca3c6926ef26bd1f6b07344aba1e91ee Mon Sep 17 00:00:00 2001 From: Bobholamovic Date: Thu, 22 Sep 2022 19:31:44 +0800 Subject: [PATCH] Allow mask not in sample in ReloadMask --- paddlers/transforms/operators.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/paddlers/transforms/operators.py b/paddlers/transforms/operators.py index a902d74..388e8df 100644 --- 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