Как определить порядок имен классов маски в загрузчике данных PyTorch? - PullRequest
0 голосов
/ 10 марта 2020

Следующий скрипт загрузчика данных считывает 11 различных имен классов из изображений 'mask'. Кажется, что индекс классов используется для определения порядка. Но если у вас есть изображение с 11 различными классами масок без назначенных имен, как первый индекс может быть «небом», а второй - «строением» и так далее? Мне трудно понять эту логику c.

class Dataset(BaseDataset):
"""CamVid Dataset. Read images, apply augmentation and preprocessing transformations.

Args:
    images_dir (str): path to images folder
    masks_dir (str): path to segmentation masks folder
    class_values (list): values of classes to extract from segmentation mask
    augmentation (albumentations.Compose): data transfromation pipeline 
        (e.g. flip, scale, etc.)
    preprocessing (albumentations.Compose): data preprocessing 
        (e.g. normalization, shape manipulation, etc.)

"""

CLASSES = ['sky', 'building', 'pole', 'road', 'pavement', 
           'tree', 'signsymbol', 'fence', 'car', 
           'pedestrian', 'bicyclist', 'unlabelled']

def __init__(
        self, 
        images_dir, 
        masks_dir, 
        classes=None, 
        augmentation=None, 
        preprocessing=None,
):
    self.ids = os.listdir(images_dir)
    self.images_fps = [os.path.join(images_dir, image_id) for image_id in self.ids]
    self.masks_fps = [os.path.join(masks_dir, image_id) for image_id in self.ids]

    # convert str names to class values on masks
    self.class_values = [self.CLASSES.index(cls.lower()) for cls in classes]

    self.augmentation = augmentation
    self.preprocessing = preprocessing

Пример изображения с масками:

enter image description here

...