<filename>datasets/cityscapes.py
import numpy as np
import torch
import torchvision.datasets as datasets
from PIL import Image

from utils.box_utils import get_bounding_box


class CityscapesDataset(datasets.Cityscapes):

    def __init__(self, root, split='train', joint_transform=None, img_transform=None):
        super(CityscapesDataset, self).__init__(root, split, target_type=['instance', 'polygon'])

        self.joint_transform = joint_transform
        self.img_transform = img_transform

    def __getitem__(self, index):
        image, target = super(CityscapesDataset, self).__getitem__(index)
        instance, json = target

        instance = self._convert_id_to_train_id(instance)
        boxes, labels = self._create_boxes(json)

        if self.joint_transform:
            image, instance, boxes, labels = self.joint_transform(image, instance, boxes, labels)

        if self.img_transform:
            image = self.img_transform(image)

        return image, instance, boxes, labels

    def _convert_id_to_train_id(self, instance):
        instance = np.array(instance)
        instance_copy = instance.copy()

        for cls in self.classes:
            instance_copy[instance == cls.id] = cls.train_id
        instance = Image.fromarray(instance_copy.astype(np.uint8))

        return instance

    def _create_boxes(self, json):
        boxes = []
        labels = []
        objects = json['objects']
        for obj in objects:
            polygons = obj['polygon']
            cls = self.get_class_from_name(obj['label'])
            if cls and cls.has_instances:
                boxes.append(get_bounding_box(polygons))
                labels.append(cls.id)

        return torch.as_tensor(boxes, dtype=torch.float32), torch.tensor(labels, dtype=torch.int64)

    @staticmethod
    def get_class_from_name(name):
        for cls in CityscapesDataset.classes:
            if cls.name == name:
                return cls
        return None

    @staticmethod
    def get_class_from_id(id):
        for cls in CityscapesDataset.classes:
            if cls.id == id:
                return cls
        return None

    @staticmethod
    def get_instance_classes():
        return [cls for cls in CityscapesDataset.classes if cls.has_instances]

    @staticmethod
    def num_instance_classes():
        return len(CityscapesDataset.get_instance_classes())

    @staticmethod
    def get_colormap():
        cmap = torch.zeros([256, 3], dtype=torch.uint8)

        for cls in CityscapesDataset.classes:
            if cls.has_instances:
                cmap[cls.trainId, :] = torch.tensor(cls.color, dtype=torch.uint8)

        return cmap

<filename>datasets/cityscapes.py
import numpy as np
import torch
import torchvision.datasets as datasets
from PIL import Image

from utils.box_utils import get_bounding_box


class CityscapesDataset(datasets.Cityscapes):

    def __init__(self, root, split='train', joint_transform=None, img_transform=None):
        super(CityscapesDataset, self).__init__(root, split, target_type=['instance', 'polygon'])

        self.joint_transform = joint_transform
        self.img_transform = img_transform

    def __getitem__(self, index):
        image, target = super(CityscapesDataset, self).__getitem__(index)
        instance, json = target

        instance = self._convert_id_to_train_id(instance)
        boxes, labels = self._create_boxes(json)

        if self.joint_transform:
            image, instance, boxes, labels = self.joint_transform(image, instance, boxes, labels)

        if self.img_transform:
            image = self.img_transform(image)

        return image, instance, boxes, labels

    def _convert_id_to_train_id(self, instance):
        instance = np.array(instance)
        instance_copy = instance.copy()

        for cls in self.classes:
            instance_copy[instance == cls.id] = cls.train_id
        instance = Image.fromarray(instance_copy.astype(np.uint8))

        return instance

    def _create_boxes(self, json):
        boxes = []
        labels = []
        objects = json['objects']
        for obj in objects:
            polygons = obj['polygon']
            cls = self.get_class_from_name(obj['label'])
            if cls and cls.has_instances:
                boxes.append(get_bounding_box(polygons))
                labels.append(cls.id)

        return torch.as_tensor(boxes, dtype=torch.float32), torch.tensor(labels, dtype=torch.int64)

    @staticmethod
    def get_class_from_name(name):
        for cls in CityscapesDataset.classes:
            if cls.name == name:
                return cls
        return None

    @staticmethod
    def get_class_from_id(id):
        for cls in CityscapesDataset.classes:
            if cls.id == id:
                return cls
        return None

    @staticmethod
    def get_instance_classes():
        return [cls for cls in CityscapesDataset.classes if cls.has_instances]

    @staticmethod
    def num_instance_classes():
        return len(CityscapesDataset.get_instance_classes())

    @staticmethod
    def get_colormap():
        cmap = torch.zeros([256, 3], dtype=torch.uint8)

        for cls in CityscapesDataset.classes:
            if cls.has_instances:
                cmap[cls.trainId, :] = torch.tensor(cls.color, dtype=torch.uint8)

        return cmap
