pytorch 训练数据以及测试 全部代码(9)---deeplab v3+ 对Cityscapes数据的处理_class cityscapessegmentation(data.dataset):

 下面是全部的代码：

import os
import torch
import numpy as np
import scipy.misc as m
from PIL import Image
from torch.utils import data
from dataloaders.utils import recursive_glob, decode_segmap
from mypath import Path
 
 
class CityscapesSegmentation(data.Dataset):
 
    def __init__(self, root=Path.db_root_dir('cityscapes'), split="train", transform=None):
 
        self.root = root
        self.split = split
        self.transform = transform
        self.files = {}
        self.n_classes = 19
 
        self.images_base = os.path.join(self.root, 'leftImg8bit', self.split)
        self.annotations_base = os.path.join(self.root, 'gtFine', self.split)
 
        self.files[split] = recursive_glob(rootdir=self.images_base, suffix='.png')
 
        self.void_classes = [0, 1, 2, 3, 4, 5, 6, 9, 10, 14, 15, 16, 18, 29, 30, -1]  # 16
        self.valid_classes = [7, 8, 11, 12, 13, 17, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 31, 32, 33]  # 19
        self.class_names = ['unlabelled', 'road', 'sidewalk', 'building', 'wall', 'fence', \
                            'pole', 'traffic_light', 'traffic_sign', 'vegetation', 'terrain', \
                            'sky', 'person', 'rider', 'car', 'truck', 'bus', 'train', \
                            'motorcycle', 'bicycle']  # 20
 
        self.ignore_index = 255
        self.class_map = dict(zip(self.valid_classes, range(self.n_classes)))
 
        if not self.files[split]:
            raise Exception("No files for split=[%s] found in %s" % (split, self.images_base))
 
        print("Found %d %s images" % (len(self.files[split]), split))
 
    def __len__(self):
        return len(self.files[self.split])
 
    def __getitem__(self, index):
 
        img_path = self.files[self.split][index].rstrip()
        lbl_path = os.path.join(self.annotations_base,
                                img_path.split(os.sep)[-2],  # os.sep=='/'  get city name
                                os.path.basename(img_path)[:-15] + 'gtFine_labelIds.png')
 
        _img = Image.open(img_path).convert('RGB')
        _tmp = np.array(Image.open(lbl_path), dtype=np.uint8)
        _tmp = self.encode_segmap(_tmp)
        _target = Image.fromarray(_tmp)
 
        sample = {'image': _img, 'label': _target}
 
        if self.transform:  # to do Data transformation or Data enhancement and  convert torch
            sample = self.transform(sample)
        return sample
 
    def encode_segmap(self, mask):  # to change original image pixel value to 0-18 and 255 according class id
        # Put all void classes to zero
        for _voidc in self.void_classes:
            mask[mask == _voidc] = self.ignore_index  # no need class and unto set 255 (white)
        for _validc in self.valid_classes:
            mask[mask == _validc] = self.class_map[_validc]  # 19 classes encode from 0 to 18
        return mask
 
 
if __name__ == '__main__':
    from dataloaders import custom_transforms as tr
    from dataloaders.utils import decode_segmap
    from torch.utils.data import DataLoader
    from torchvision import transforms
    import matplotlib.pyplot as plt  # to show image
 
    composed_transforms_tr = transforms.Compose([
        tr.RandomHorizontalFlip(),
        tr.RandomScale((0.5, 0.75)),
        tr.RandomCrop((512, 1024)),
        tr.RandomRotate(5),
        tr.ToTensor()])
 
    cityscapes_train = CityscapesSegmentation(split='train',
                                transform=composed_transforms_tr)
    dataloader = DataLoader(cityscapes_train, batch_size=2, shuffle=True, num_workers=2)
 
    for ii, sample in enumerate(dataloader):
        for jj in range(sample["image"].size()[0]):
            img = sample['image'].numpy()  # from torch convert to numpy n x c x h x w
            gt = sample['label'].numpy()  # from torch convert to numpy n x c x h x w
            tmp = np.array(gt[jj]).astype(np.uint8)  # tmp.shape=c x h x w
            tmp = np.squeeze(tmp, axis=0)  # if c=1,tmp.shape=c x h x w; or tmp.shape=c x h x w
            segmap = decode_segmap(tmp, dataset='cityscapes')
            img_tmp = np.transpose(img[jj], axes=[1, 2, 0]).astype(np.uint8)  # img_tmp=h x w x c
            plt.figure()
            plt.title('display')
            plt.subplot(211)
            plt.imshow(img_tmp)
            plt.subplot(212)
            plt.imshow(segmap)
 
        if ii == 1:
            break
    plt.show(block=True)
 

下面怎么读取图片的 可以参考：https://blog.csdn.net/zz2230633069/article/details/84640867

self.files[split] = recursive_glob(rootdir=self.images_base, suffix='.png')

转换的为：

composed_transforms_tr = transforms.Compose([
        tr.RandomHorizontalFlip(),
        tr.RandomScale((0.5, 0.75)),
        tr.RandomCrop((512, 1024)),
        tr.RandomRotate(5),
        tr.ToTensor()])

上面关于图像变换或者说增强的实现代码如下：

上面的前四个变换都保持了原图和标签的type为PIL.PngImagePlugin.PngImageFile，这些图的像素数值大小和类型（uint8）不发生改变，结构也没有变化（原图为h x w x 3，标签图为h x w)

class RandomHorizontalFlip(object):
    def __call__(self, sample):
        img = sample['image']
        mask = sample['label']
        if random.random() < 0.5:
            img = img.transpose(Image.FLIP_LEFT_RIGHT)
            mask = mask.transpose(Image.FLIP_LEFT_RIGHT)
 
        return {'image': img,
                'label': mask}
 
 
class RandomScale(object):
    def __init__(self, limit):
        self.limit = limit
 
    def __call__(self, sample):
        img = sample['image']
        mask = sample['label']
        assert img.size == mask.size
 
        scale = random.uniform(self.limit[0], self.limit[1])
        w = int(scale * img.size[0])
        h = int(scale * img.size[1])
 
        img, mask = img.resize((w, h), Image.BILINEAR), mask.resize((w, h), Image.NEAREST)
 
        return {'image': img, 'label': mask}
 
 
class RandomCrop(object):
    def __init__(self, size, padding=0):
        if isinstance(size, numbers.Number):
            self.size = (int(size), int(size))
        else:
            self.size = size # h, w
        self.padding = padding
 
    def __call__(self, sample):
        img, mask = sample['image'], sample['label']
 
        if self.padding > 0:
            img = ImageOps.expand(img, border=self.padding, fill=0)
            mask = ImageOps.expand(mask, border=self.padding, fill=0)
 
        assert img.size == mask.size
        w, h = img.size
        th, tw = self.size # target size
        if w == tw and h == th:
            return {'image': img,
                    'label': mask}
        if w < tw or h < th:
            img = img.resize((tw, th), Image.BILINEAR)
            mask = mask.resize((tw, th), Image.NEAREST)
            return {'image': img,
                    'label': mask}
 
        x1 = random.randint(0, w - tw)
        y1 = random.randint(0, h - th)
        img = img.crop((x1, y1, x1 + tw, y1 + th))
        mask = mask.crop((x1, y1, x1 + tw, y1 + th))
 
        return {'image': img,
                'label': mask}
 
 
class RandomRotate(object):
    def __init__(self, degree):
        self.degree = degree
 
    def __call__(self, sample):
        img = sample['image']
        mask = sample['label']
        rotate_degree = random.random() * 2 * self.degree - self.degree
        img = img.rotate(rotate_degree, Image.BILINEAR)
        mask = mask.rotate(rotate_degree, Image.NEAREST)
 
        return {'image': img,
                'label': mask}
 
class ToTensor(object):
    """Convert ndarrays in sample to Tensors."""
 
    def __call__(self, sample):
        # swap color axis because
        # numpy image: H x W x C
        # torch image: C X H X W
        img = np.array(sample['image']).astype(np.float32).transpose((2, 0, 1))
        mask = np.expand_dims(np.array(sample['label']).astype(np.float32), -1).transpose((2, 0, 1))
        mask[mask == 255] = 0   #
 
        img = torch.from_numpy(img).float()
        mask = torch.from_numpy(mask).float()
 
 
        return {'image': img,
                'label': mask}

直到第五个也就是最后一个（ToTensor函数）变化，对原图首先从PIL.PngImagePlugin.PngImageFile变到numpy类型同时数据类型从uint8变为float32类型，然后维度变化从（h x w x c ）到（c x h x w），最后从numpy类型变为torch的tensor类型，同时强制将数据类型为torch.FloatTensor。这样，就将原图转变为一个tensor可以输入后面的深度学习网络中了。

与此相对的标签图也是从PIL.PngImagePlugin.PngImageFile变到numpy类型同时数据类型从uint8变为float32类型，然后维度变化从（h x w  ）增加一维得到（h x w x 1）接着调整维度到（1 x h x w），然后mask里面的数值进行处理：255.值大小的全部被重置为0，所以mask里面的值现在只有0-18这些数字了；最后从numpy类型变为torch的tensor类型，同时强制将数据类型为torch.FloatTensor。这样，就将标签图转变为一个tensor可以输入后面的深度学习网络中了。

对上面的两个tensor的重新变成图像的代码如下:

    for ii, sample in enumerate(dataloader):
        for jj in range(sample["image"].size()[0]):
            img = sample['image'].numpy()  # from torch convert to numpy n x 3 x h x w
            gt = sample['label'].numpy()  # from torch convert to numpy n x 1 x h x w
            tmp = np.array(gt[jj]).astype(np.uint8)  # tmp.shape=1 x h x w
            tmp = np.squeeze(tmp, axis=0)  # if c=1,tmp.shape=h x w; or tmp.shape=c x h x w dimension-reduction
            segmap = decode_segmap(tmp, dataset='cityscapes') 
            img_tmp = np.transpose(img[jj], axes=[1, 2, 0]).astype(np.uint8)  # img_tmp=h x w x 3
            plt.figure()
            plt.title('display')
            plt.subplot(211)
            plt.imshow(img_tmp)
            plt.subplot(212)
            plt.imshow(segmap)
 
        if ii == 1:
            break
    plt.show(block=True)

里面的标签图（h x w）解码代码如下：

只要是同一类的就给相应的RGB数值，然后整合三张图到一张图里面

segmap = decode_segmap(tmp, dataset='cityscapes')  # tmp.shape=h x w
def decode_segmap(label_mask, dataset, plot=False):
    """Decode segmentation class labels into a color image
    Args:
        label_mask (np.ndarray): an (M,N) array of integer values denoting
          the class label at each spatial location.
        plot (bool, optional): whether to show the resulting color image
          in a figure.
    Returns:
        (np.ndarray, optional): the resulting decoded color image.
    """
    if dataset == 'pascal':
        n_classes = 21
        label_colours = get_pascal_labels()
    elif dataset == 'cityscapes':
        n_classes = 19
        label_colours = get_cityscapes_labels()
    else:
        raise NotImplementedError
 
    r = label_mask.copy()  # h x w
    g = label_mask.copy()  # h x w
    b = label_mask.copy()  # h x w
    for ll in range(0, n_classes):
        r[label_mask == ll] = label_colours[ll, 0]
        g[label_mask == ll] = label_colours[ll, 1]
        b[label_mask == ll] = label_colours[ll, 2]
    rgb = np.zeros((label_mask.shape[0], label_mask.shape[1], 3)) # h x w x 3初始化
    rgb[:, :, 0] = r / 255.0
    rgb[:, :, 1] = g / 255.0
    rgb[:, :, 2] = b / 255.0
    if plot:
        plt.imshow(rgb)
        plt.show()
    else:
        return rgb

下面就是label_colours的和类别对应色彩代码详情可以看cityscapes的标签颜色对照表https://blog.csdn.net/zz2230633069/article/details/84591532：

def get_cityscapes_labels():
    return np.array([
        # [  0,   0,   0],
        [128, 64, 128],
        [244, 35, 232],
        [70, 70, 70],
        [102, 102, 156],
        [190, 153, 153],
        [153, 153, 153],
        [250, 170, 30],
        [220, 220, 0],
        [107, 142, 35],
        [152, 251, 152],
        [0, 130, 180],
        [220, 20, 60],
        [255, 0, 0],
        [0, 0, 142],
        [0, 0, 70],
        [0, 60, 100],
        [0, 80, 100],
        [0, 0, 230],
        [119, 11, 32]])
def get_pascal_labels():
    """Load the mapping that associates pascal classes with label colors
    Returns:
        np.ndarray with dimensions (21, 3)
    """
    return np.asarray([[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]])