From a0973a382ce785bb6aa8d3efa0f187f000b90ee5 Mon Sep 17 00:00:00 2001 From: zhangpu <1250681871@qq.com> Date: Tue, 9 Dec 2025 18:33:18 +0800 Subject: [PATCH] =?UTF-8?q?=E5=A4=84=E7=90=86=E7=B1=BB=E5=9E=8B=E5=BC=82?= =?UTF-8?q?=E5=B8=B8=E3=80=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- SegFormer/training/dataset.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/SegFormer/training/dataset.py b/SegFormer/training/dataset.py index a71340d..de782e0 100644 --- a/SegFormer/training/dataset.py +++ b/SegFormer/training/dataset.py @@ -161,10 +161,13 @@ class FoodSegmentationDataset(Dataset): # albumentations会自动将增强同时应用到image和mask transformed = self.transform(image=image, mask=mask) image = transformed['image'] # Tensor (3, H, W) - mask = transformed['mask'] # ndarray (H, W) + mask = transformed['mask'] # 可能是ndarray或Tensor - # 3. 将mask转换为Tensor - mask = torch.from_numpy(mask).long() + # 3. 将mask转换为Tensor（检查类型） + if isinstance(mask, np.ndarray): + mask = torch.from_numpy(mask).long() + else: + mask = mask.long() # 已经是Tensor，直接转换类型 # 4. 检查mask的有效性 # 确保mask的值在[0, num_classes-1]范围内