99 lines
4.2 KiB
Python
99 lines
4.2 KiB
Python
import os
|
||
import cv2
|
||
import numpy as np
|
||
import torch
|
||
from torch.utils.data import Dataset
|
||
|
||
class SegmentationDataset(Dataset):
|
||
"""
|
||
自定义图像分割数据集。
|
||
根据指定的 seg_mode，可以生成适用于 'multiclass' 或 'multilabel' 任务的掩码。
|
||
"""
|
||
def __init__(self, image_dir, mask_dir, classes, class_rgb_values, seg_mode, augmentation=None):
|
||
"""
|
||
Args:
|
||
image_dir (str): 图像文件目录。
|
||
mask_dir (str): 掩码文件目录。
|
||
classes (list): 类别名称列表。
|
||
class_rgb_values (list): 每个类别在灰度掩码中对应的像素值。
|
||
augmentation (albumentations.Compose, optional): 数据增强流程。
|
||
seg_mode (str): 分割模式, 'multiclass' 或 'multilabel'。
|
||
"""
|
||
self.image_dir = image_dir
|
||
self.mask_dir = mask_dir
|
||
self.image_filenames = sorted(os.listdir(image_dir))
|
||
self.augmentation = augmentation
|
||
self.classes = classes
|
||
self.class_rgb_values = class_rgb_values
|
||
|
||
# 【新增】存储分割模式并进行验证
|
||
self.seg_mode = seg_mode
|
||
if self.seg_mode not in ['multiclass', 'multilabel']:
|
||
raise ValueError(f"seg_mode must be 'multiclass' or 'multilabel', but got {self.seg_mode}")
|
||
|
||
print(f"Found {len(self.image_filenames)} images in {image_dir}. Dataset mode: '{self.seg_mode}'")
|
||
|
||
def __len__(self):
|
||
return len(self.image_filenames)
|
||
|
||
def __getitem__(self, idx):
|
||
img_path = os.path.join(self.image_dir, self.image_filenames[idx])
|
||
image = cv2.imread(img_path)
|
||
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
||
|
||
mask_path = os.path.join(self.mask_dir, self.image_filenames[idx])
|
||
# 读取灰度图
|
||
mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
|
||
|
||
if mask is None:
|
||
raise FileNotFoundError(f"Mask file not found or could not be read: {mask_path}")
|
||
|
||
# 【核心修改】根据 seg_mode 选择掩码的处理方式
|
||
if self.seg_mode == 'multilabel':
|
||
# 为 'multilabel' 模式创建 one-hot 编码掩码
|
||
processed_mask = self._create_one_hot_mask(mask)
|
||
else: # 'multiclass'
|
||
# 为 'multiclass' 模式创建类索引掩码
|
||
processed_mask = self._create_class_index_mask(mask)
|
||
|
||
if self.augmentation:
|
||
# Albumentations 可以同时处理 (H,W,C) 和 (H,W) 格式的掩码
|
||
sample = self.augmentation(image=image, mask=processed_mask)
|
||
image = sample['image']
|
||
mask = sample['mask']
|
||
else:
|
||
mask = processed_mask
|
||
|
||
# 【核心修改】根据 seg_mode 对最终的掩码张量进行处理
|
||
if self.seg_mode == 'multilabel':
|
||
# 调整维度顺序 (H, W, C) -> (C, H, W) 并转换为 float
|
||
mask = mask.permute(2, 0, 1).float()
|
||
else: # 'multiclass'
|
||
# 直接转换为 long 类型，不需要调整维度
|
||
mask = mask.long()
|
||
|
||
return image, mask
|
||
|
||
def _create_one_hot_mask(self, mask):
|
||
"""
|
||
将单通道的灰度掩码 (H, W) 转换为 one-hot 编码的掩码 (H, W, C)。
|
||
这是为 'multilabel' 模式准备的。
|
||
返回一个 NumPy 数组。
|
||
"""
|
||
semantic_map = np.zeros((mask.shape[0], mask.shape[1], len(self.class_rgb_values)), dtype=np.uint8)
|
||
for i, value in enumerate(self.class_rgb_values):
|
||
semantic_map[:, :, i] = (mask == value).astype(np.uint8)
|
||
return semantic_map
|
||
|
||
def _create_class_index_mask(self, mask):
|
||
"""
|
||
【新增函数】
|
||
将单通道的灰度掩码 (H, W) 转换为类索引掩码 (H, W)。
|
||
这是为 'multiclass' 模式准备的。
|
||
返回一个 NumPy 数组。
|
||
"""
|
||
class_index_map = np.zeros(mask.shape, dtype=np.uint8)
|
||
for i, value in enumerate(self.class_rgb_values):
|
||
# 将灰度值为 value 的像素，其类别索引设置为 i
|
||
class_index_map[mask == value] = i
|
||
return class_index_map |