处理类型异常。

2025-12-09 18:33:18 +08:00
parent 56de053db1
commit a0973a382c
1 changed files with 6 additions and 3 deletions
@@ -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]范围内