From e547a8e0b8dceed3d9f49b09a64252edd214c5b8 Mon Sep 17 00:00:00 2001 From: chflame Date: Sun, 19 May 2024 19:22:35 +0800 Subject: [PATCH] fix bug of CorpByMaskV2 --- py/crop_by_mask_v2.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/py/crop_by_mask_v2.py b/py/crop_by_mask_v2.py index 7b7d8eb..f1986db 100644 --- a/py/crop_by_mask_v2.py +++ b/py/crop_by_mask_v2.py @@ -45,11 +45,11 @@ class CropByMaskV2: for l in image: l_images.append(torch.unsqueeze(l, 0)) if mask.dim() == 2: - mask = torch.unsqueeze(mask_for_crop, 0) + mask = torch.unsqueeze(mask, 0) # 如果有多张mask输入，使用第一张 if mask.shape[0] > 1: log(f"Warning: Multiple mask inputs, using the first.", message_type='warning') - mask = torch.unsqueeze(mask_for_crop[0], 0) + mask = torch.unsqueeze(mask[0], 0) if invert_mask: mask = 1 - mask l_masks.append(tensor2pil(torch.unsqueeze(mask, 0)).convert('L'))