from enum import Enum

import torch


class AnnotationType(Enum):
    """Enum for possible annotation types

    Each dataset in this library comes with an attribute which defines which
    annotations are supplied with the dataset.
    """
    NONE = 0
    OBJECT_DETECTION_2D = 1
    OBJECT_DETECTION_3D = 2
    SEMANTIC_SEGMENTATION = 3
    INSTANCE_SEGMENTATION = 4


def collate_object_detection_2d(batch):
    imgs = torch.stack([elem[0] for elem in batch])
