d2l图像分割数据集制作

我想吃鱼了，

墨衍会员 · AI 创作全网分发码龄5年

关注

原创 2023-04-06 14:39:25 637 阅读 AI 写同款· GEO 优化 ›

动手学习深度学习14-微调-热狗识别代码图像识别，微调阅读详情

本文主要讲述图像分割Segmentation数据集的制作

加载：

实现从RGB三个通道到类别索引单int值：

图像增广

集成用法

加载：

在Segmentation里面有三个txt，里面写的是JPEG文件夹里面的图片的名称
将图片转为RGB，然后提取images的图片名，读取对应文件中的图片作为features
这里注意，对于语义分割，其label是对每个像素进行标注，所以标注可以为一张与features等大小的图片，是png文件是为了防止jpg图片格式会压缩：

def read_voc_images(voc_dir, is_train=True):
    """读取所有VOC图像并标注"""
    txt_fname = os.path.join(voc_dir, 'ImageSets', 'Segmentation',
                            'train.txt' if is_train else 'val.txt')
    mode = torchvision.io.image.ImageReadMode.RGB
    with open(txt_fname, 'r') as f:
        images = f.read().split()
    features, labels = [], []
    for i, fname in enumerate(images):
        features.append(torchvision.io.read_image(os.path.join(
            voc_dir, 'JPEGImages', f'{fname}.jpg')))
        labels.append(torchvision.io.read_image(os.path.join(
            voc_dir, 'SegmentationClass' ,f'{fname}.png'), mode))
    return features, labels

voc_dir = '/CV/xhr/VOCdevkit/VOC2012'

train_features, train_labels = read_voc_images(voc_dir, True)
len(train_features), train_features[1].shape, train_labels[1].shape

'''
(1464, torch.Size([3, 375, 500]), torch.Size([3, 375, 500]))
'''

由此可见，该函数得到的是list，里面的每个元素对应的是每个图片的tensor(c,w,h)

显示一下：

n = 5
imgs = train_features[0:n] + train_labels[0:n]
imgs = [img.permute(1,2,0) for img in imgs]
d2l.show_images(imgs, 2, n);

对于Segmentation中的标记，为对各个像素的标记，每个类别对应不同的颜色，具体颜色与label标签对应见下：

VOC_COLORMAP = [[0, 0, 0], [128, 0, 0], [0, 128, 0], [128, 128, 0],
                [0, 0, 128], [128, 0, 128], [0, 128, 128], [128, 128, 128],
                [64, 0, 0], [192, 0, 0], [64, 128, 0], [192, 128, 0],
                [64, 0, 128], [192, 0, 128], [64, 128, 128], [192, 128, 128],
                [0, 64, 0], [128, 64, 0], [0, 192, 0], [128, 192, 0],
                [0, 64, 128]]
#@save
VOC_CLASSES = ['background', 'aeroplane', 'bicycle', 'bird', 'boat',
                'bottle', 'bus', 'car', 'cat', 'chair', 'cow',
                'diningtable', 'dog', 'horse', 'motorbike', 'person',
                'potted plant', 'sheep', 'sofa', 'train', 'tv/monitor']

这里返回每个label的RGB三个通道情况，第一张图片飞机用红色标注，所以只有在第一个通道R上有数值128，其余的两个通道对应位置全为0

train_labels[0][0][50:150, 230:340], train_labels[0][1][50:150, 230:340]

'''
(tensor([[  0,   0,   0,  ...,   0,   0,   0],
         [  0,   0,   0,  ...,   0,   0,   0],
         [  0,   0,   0,  ...,   0,   0,   0],
         ...,
         [128, 128, 128,  ...,   0,   0,   0],
         [128, 128, 128,  ...,   0,   0,   0],
         [128, 128, 128,  ...,   0,   0,   0]], dtype=torch.uint8),
 tensor([[0, 0, 0,  ..., 0, 0, 0],
         [0, 0, 0,  ..., 0, 0, 0],
         [0, 0, 0,  ..., 0, 0, 0],
         ...,
         [0, 0, 0,  ..., 0, 0, 0],
         [0, 0, 0,  ..., 0, 0, 0],
         [0, 0, 0,  ..., 0, 0, 0]], dtype=torch.uint8))
'''

实现从RGB三个通道到类别索引单int值：

目的：在label中，使用不同颜色进行标注，但是每个像素返回的是一个rgb三通道list，不好训练，所以要将不同的rgb通道弄成一个int，方便训练.--hash索引(longtensor)
3个256的RGB通道共有256^3种组合，所以先创建一个大的索引tensor,再使用一种将不同通道上的像素进行不同运算的方法,即(r×256+g)×256+b，使上述标记种类的rgb都能够索引到相应的int值。
返回的这个colormap是一个long.tensor，里面对应上面种类的21个不同位置处不是默认值，是从0-20的类别对应数目，其他的全是默认创建值。
下面这个例子表示了索引张量tensor中对应位置对应元素的过程。在label中，相应的颜色映射到对应的整数i中。

def voc_colormap2label():
    """构建从RGB到VOC类别索引的映射"""
    colormap2label = torch.zeros(256 ** 3, dtype=torch.long)
    for i, colormap in enumerate(VOC_COLORMAP):
        colormap2label[
            (colormap[0] * 256 + colormap[1]) * 256 + colormap[2]] = i
    return colormap2label

def voc_label_indices(colormap, colormap2label):
    """将VOC标签中的RGB值映射到它们的类别索引"""
    colormap = colormap.permute(1, 2, 0).numpy().astype('int32')
    idx = ((colormap[:, :, 0] * 256 + colormap[:, :, 1]) * 256
        + colormap[:, :, 2])
    return colormap2label[idx]

输入对应的图片(第一张),再输入映射函数，则返回的每个像素为对应类的对应编号：

y = voc_label_indices(train_labels[0], voc_colormap2label())
y[50:150, 230:340], VOC_CLASSES[20]

'''
(tensor([[0, 0, 0,  ..., 0, 0, 0],
         [0, 0, 0,  ..., 0, 0, 0],
         [0, 0, 0,  ..., 0, 0, 0],
         ...,
         [1, 1, 1,  ..., 0, 0, 0],
         [1, 1, 1,  ..., 0, 0, 0],
         [1, 1, 1,  ..., 0, 0, 0]]),
 'tv/monitor')
'''

注意，索引用到了hash表：

在colormap2label中，是一个1维tensor，此时idx为2维tensor，每个元素显示的为通道数上元素的对应计算值，此时传入这个1维tensor中，会输出一个相同维度的tensor，里面的每个元素为对应‘字典’中对应计算值的索引值‘i'。

对于没标记成类别的颜色，还有全黑的背景，全部对应为background类，在colormap2label的'字典'中索引的全是0。

如这里的14737...是对应标记白边：

图像增广

在对图像进行随机裁剪时，同时也要对label进行裁剪
torchvision.transforms.RandomCrop.get_params命令传入图片和裁剪的尺寸，返回的是一个4元素元组，分别为裁剪区域的左上角xy坐标与裁剪wh，可以进行functional.crop传入图像与该元组，即对指定裁剪区域进行裁剪。

def voc_rand_crop(feature, label, height, width):
    """随机裁剪特征和标签图像"""
    rect = torchvision.transforms.RandomCrop.get_params(
        feature, (height, width))
    feature = torchvision.transforms.functional.crop(feature, *rect)
    label = torchvision.transforms.functional.crop(label, *rect)
    return feature, label

segmentation不好做resize，因为resize是对像素值插值，插值对应不了label。

一般都用crop

filter是假设图像的尺寸如果小于crop的尺寸，就直接去掉了。

getitem对应索引得到的是经过crop等teansform的图片！
在这里注意，features是通过read_voc_images得到的。对应的是一个list，里面每个元素才为对应的图片cwh，这里self.features是经过了filter的，在filter函数中，返回的img是for img in imgs，其中这个imgs是传入的以每个元素为图片的list，这样返回的feature就是具体的图片了，self.labels也是一样的道理！！

class VOCSegDataset(torch.utils.data.Dataset):
    """⼀个⽤于加载VOC数据集的⾃定义数据集"""
    def __init__(self, is_train, crop_size, voc_dir):
        self.transform = torchvision.transforms.Normalize(
            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        self.crop_size = crop_size
        features, labels = read_voc_images(voc_dir, is_train=is_train) # 这里是list
        self.features = [self.normalize_image(feature)
                        for feature in self.filter(features)]
        self.labels = self.filter(labels)
        self.colormap2label = voc_colormap2label()
        print('read ' + str(len(self.features)) + ' examples')
    def normalize_image(self, img):
        return self.transform(img.float() / 255)
    
    def filter(self, imgs): 
        # imgs是每个元素为tensor(c,w,h)的list，for循环后返回的即为img的tensor
        return [img for img in imgs if (
        img.shape[1] >= self.crop_size[0] and
        img.shape[2] >= self.crop_size[1])]
    
    def __getitem__(self, idx):
        feature, label = voc_rand_crop(self.features[idx], self.labels[idx],
                                        *self.crop_size)
        return (feature, voc_label_indices(label, self.colormap2label))
    
    def __len__(self):
        return len(self.features)

crop_size = (320, 480)
voc_train = VOCSegDataset(True, crop_size, voc_dir)
voc_test = VOCSegDataset(False, crop_size, voc_dir)

'''
read 1114 examples
read 1078 examples
'''

注意：label里面没有c，对应的是之前提到的将rgb通道计算为单个int值！

batch_size = 64
train_iter = torch.utils.data.DataLoader(voc_train, batch_size, shuffle=True,
                                    drop_last=True,
                                    num_workers=d2l.get_dataloader_workers())
for X, Y in train_iter:
    print(X.shape)
    print(Y.shape)
    break

'''
torch.Size([64, 3, 320, 480])
torch.Size([64, 320, 480])
'''

集成用法

def load_data_voc(batch_size, crop_size):
    """加载VOC语义分割数据集"""
    voc_dir = '/CV/xhr/VOCdevkit/VOC2012'
    num_workers = d2l.get_dataloader_workers()
    train_iter = torch.utils.data.DataLoader(
        VOCSegDataset(True, crop_size, voc_dir), batch_size,
        shuffle=True, drop_last=True, num_workers=num_workers)
    test_iter = torch.utils.data.DataLoader(
        VOCSegDataset(False, crop_size, voc_dir), batch_size,
        drop_last=True, num_workers=num_workers)
    return train_iter, test_iter