from __future__ import annotations
from enum import Enum as BaseEnum
from enum import EnumMeta as BaseEnumMeta
from typing import Any, Dict, List, Optional, Tuple, Union
from pydantic import BaseModel as PydanticBaseModel
from pydantic import ConfigDict, field_validator
from segments.exceptions import ValidationError
from typing_extensions import Literal, TypedDict
class BaseModel(PydanticBaseModel):
model_config = ConfigDict(
# What happens with extra fields in dictionaries. Use ignore in production and allow in debug mode. https://pydantic-docs.helpmanual.io/usage/model_config/#change-behaviour-globally
extra="ignore",
# What happens with wrong field types. Use false in production and true in debug mode. https://pydantic-docs.helpmanual.io/usage/types/#arbitrary-types-allowed
arbitrary_types_allowed=False,
)
class EnumMeta(BaseEnumMeta):
# https://stackoverflow.com/questions/43634618/how-do-i-test-if-int-value-exists-in-python-enum-without-using-try-catch
def __contains__(self, item):
return isinstance(item, self) or item in {
v.value for v in self.__members__.values()
}
# https://stackoverflow.com/questions/29503339/how-to-get-all-values-from-python-enum-class
def __str__(self):
return ", ".join(c.value for c in self)
def __repr__(self):
return self.__str__()
class Enum(BaseEnum, metaclass=EnumMeta):
pass
#################################
# Literals, constants and enums #
#################################
[docs]class LabelStatus(str, Enum):
REVIEWED = "REVIEWED"
REVIEWING_IN_PROGRESS = "REVIEWING_IN_PROGRESS"
LABELED = "LABELED"
LABELING_IN_PROGRESS = "LABELING_IN_PROGRESS"
REJECTED = "REJECTED"
PRELABELED = "PRELABELED"
SKIPPED = "SKIPPED"
UNLABELED = "UNLABELED"
VERIFIED = "VERIFIED"
[docs]class TaskType(str, Enum):
SEGMENTATION_BITMAP = "segmentation-bitmap"
SEGMENTATION_BITMAP_HIGHRES = "segmentation-bitmap-highres"
IMAGE_VECTOR_SEQUENCE = "image-vector-sequence"
IMAGE_SEGMENTATION_SEQUENCE = "image-segmentation-sequence"
BBOXES = "bboxes"
VECTOR = "vector"
KEYPOINTS = "keypoints"
POINTCLOUD_CUBOID = "pointcloud-cuboid"
POINTCLOUD_CUBOID_SEQUENCE = "pointcloud-cuboid-sequence"
POINTCLOUD_SEGMENTATION = "pointcloud-segmentation"
POINTCLOUD_SEGMENTATION_SEQUENCE = "pointcloud-segmentation-sequence"
POINTCLOUD_VECTOR = "pointcloud-vector"
POINTCLOUD_VECTOR_SEQUENCE = "pointcloud-vector-sequence"
MULTISENSOR = "multisensor"
MULTISENSOR_SEQUENCE = "multisensor-sequence"
TEXT_NAMED_ENTITIES = "text-named-entities"
TEXT_SPAN_CATEGORIZATION = "text-span-categorization"
EMPTY = ""
class Role(str, Enum):
LABELER = "labeler"
REVIEWER = "reviewer"
MANAGER = "manager"
ADMIN = "admin"
[docs]class IssueStatus(str, Enum):
OPEN = "OPEN"
CLOSED = "CLOSED"
class Category(str, Enum):
STREET_SCENERY = "street_scenery"
GARDEN = "garden"
AGRICULTURE = "agriculture"
SATELLITE = "satellite"
PEOPLE = "people"
MEDICAL = "medical"
FRUIT = "fruit"
OTHER = "other"
[docs]class CameraConvention(str, Enum):
OPEN_CV = "OpenCV"
OPEN_GL = "OpenGL"
class InputType(str, Enum):
SELECT = "select"
TEXT = "text"
NUMBER = "number"
CHECKBOX = "checkbox"
[docs]class CameraDistortionModel(str, Enum):
FISH_EYE = "fisheye"
BROWN_CONRADY = "brown-conrady"
[docs]class PCDType(str, Enum):
PCD = "pcd"
KITTI = "kitti"
NUSCENES = "nuscenes"
[docs]class ReleaseStatus(str, Enum):
PENDING = "PENDING"
SUCCEEDED = "SUCCEEDED"
FAILED = "FAILED"
[docs]class ImageVectorAnnotationType(str, Enum):
BBOX = "bbox"
POLYGON = "polygon"
POLYLINE = "polyline"
POINT = "point"
[docs]class PointcloudCuboidAnnotationType(str, Enum):
CUBOID = "cuboid"
CUBOID_SYNC = "cuboid-sync"
[docs]class PointcloudVectorAnnotationType(str, Enum):
POLYGON = "polygon"
POLYLINE = "polyline"
POINT = "point"
class ExportFormat(str, Enum):
COCO_PANOPTIC = "coco-panoptic"
COCO_INSTANCE = "coco-instance"
YOLO = "yolo"
INSTANCE = "instance"
INSTANCE_COLOR = "instance-color"
SEMANTIC = "semantic"
SEMANTIC_COLOR = "semantic-color"
POLYGON = "polygon"
RGB = Tuple[int, int, int]
RGBA = Tuple[int, int, int, int]
FormatVersion = Union[float, str]
ObjectAttributes = Dict[str, Optional[Union[str, bool]]]
ImageAttributes = Dict[str, Optional[Union[str, bool]]]
###########
# Release #
###########
[docs]class URL(BaseModel):
url: Optional[
str
] = None # TODO Remove optional (e.g., the backend does not return an URL when adding a release).
[docs]class Release(BaseModel):
uuid: str
name: str
description: str
release_type: Literal["JSON"]
attributes: URL
status: ReleaseStatus
# status_info: str
created_at: str
samples_count: int
#########
# Issue #
#########
[docs]class Issue(BaseModel):
uuid: str
description: str
created_at: str
updated_at: str
created_by: str
updated_by: str
comments: List[IssueComment]
status: IssueStatus
sample_uuid: str
sample_name: str
########
# File #
########
# https://stackoverflow.com/questions/60003444/typeddict-when-keys-have-invalid-names
AWSFields = TypedDict(
"AWSFields",
{
"acl": str,
"Content-Type": str,
"key": str,
"x-amz-algorithm": str,
"x-amz-credential": str,
"x-amz-date": str,
"policy": str,
"x-amz-signature": str,
},
)
[docs]class PresignedPostFields(BaseModel):
url: str
fields: AWSFields
[docs]class File(BaseModel):
uuid: str
filename: str
url: str
created_at: str
presignedPostFields: PresignedPostFields
#########
# Label #
#########
[docs]class Annotation(BaseModel):
id: int
category_id: int
attributes: Optional[ObjectAttributes] = None
# Image segmentation
[docs]class ImageSegmentationLabelAttributes(BaseModel):
annotations: List[Annotation]
segmentation_bitmap: URL
image_attributes: Optional[ImageAttributes] = None
format_version: Optional[FormatVersion] = None
# Image vector
# https://stackoverflow.com/questions/51575931/class-inheritance-in-python-3-7-dataclasses
[docs]class ImageVectorAnnotation(BaseModel):
id: int
category_id: int
points: List[List[float]]
type: ImageVectorAnnotationType
attributes: Optional[ObjectAttributes] = None
[docs]class ImageVectorLabelAttributes(BaseModel):
annotations: List[ImageVectorAnnotation]
format_version: Optional[FormatVersion] = None
image_attributes: Optional[ImageAttributes] = None
# Image sequence segmentation
class ImageSequenceSegmentationAnnotation(Annotation):
track_id: int
is_keyframe: bool = False
class ImageSequenceSegmentationFrame(ImageSegmentationLabelAttributes):
annotations: List[ImageSequenceSegmentationAnnotation]
timestamp: Optional[Union[str, int]]
format_version: Optional[FormatVersion]
class ImageSequenceSegmentationLabelAttributes(BaseModel):
frames: List[ImageSequenceSegmentationFrame]
format_version: Optional[FormatVersion]
# Image sequence vector
[docs]class ImageSequenceVectorAnnotation(ImageVectorAnnotation):
track_id: int
is_keyframe: bool = False
[docs]class ImageVectorFrame(ImageVectorLabelAttributes):
annotations: List[ImageSequenceVectorAnnotation]
timestamp: Optional[Union[str, int]] = None
format_version: Optional[FormatVersion] = None
image_attributes: Optional[ImageAttributes] = None
[docs]class ImageSequenceVectorLabelAttributes(BaseModel):
frames: List[ImageVectorFrame]
format_version: Optional[FormatVersion] = None
# Point cloud segmentation
[docs]class PointcloudSegmentationLabelAttributes(BaseModel):
annotations: List[Annotation]
point_annotations: List[int]
format_version: Optional[FormatVersion] = None
[docs]class XYZ(BaseModel):
x: float
y: float
z: float
[docs]class XYZW(BaseModel):
qx: float
qy: float
qz: float
qw: float
[docs]class FisheyeDistortionCoefficients(BaseModel):
k1: float
k2: float
k3: float
k4: float
[docs]class BrownConradyDistortionCoefficients(BaseModel):
k1: float
k2: float
p1: float
p2: float
k3: float
[docs]class Distortion(BaseModel):
model: CameraDistortionModel
coefficients: Union[
FisheyeDistortionCoefficients, BrownConradyDistortionCoefficients
]
# Point cloud cuboid
# https://stackoverflow.com/questions/51575931/class-inheritance-in-python-3-7-dataclasses
[docs]class PointcloudCuboidAnnotation(BaseModel):
id: int
category_id: int
position: XYZ
dimensions: XYZ
yaw: float
rotation: Optional[XYZW] = None
type: PointcloudCuboidAnnotationType
attributes: Optional[ObjectAttributes] = None
[docs]class PointcloudCuboidLabelAttributes(BaseModel):
annotations: List[PointcloudCuboidAnnotation]
format_version: Optional[FormatVersion] = None
# Point cloud vector
[docs]class PointcloudVectorAnnotation(BaseModel):
id: int
category_id: int
points: List[List[float]]
type: PointcloudVectorAnnotationType
attributes: Optional[ObjectAttributes] = None
[docs]class PointcloudVectorLabelAttributes(BaseModel):
annotations: List[PointcloudVectorAnnotation]
format_version: Optional[FormatVersion] = None
# Point cloud sequence segmentation
[docs]class PointcloudSequenceSegmentationAnnotation(Annotation):
track_id: int
is_keyframe: bool = False
attributes: Optional[ObjectAttributes] = None
[docs]class PointcloudSegmentationFrame(PointcloudSegmentationLabelAttributes):
annotations: List[PointcloudSequenceSegmentationAnnotation]
point_annotations: List[int]
timestamp: Optional[Union[str, int]] = None
format_version: Optional[FormatVersion] = None
[docs]class PointcloudSequenceSegmentationLabelAttributes(BaseModel):
frames: List[PointcloudSegmentationFrame]
format_version: Optional[FormatVersion] = None
# Point cloud sequence cuboid
[docs]class PointcloudSequenceCuboidAnnotation(PointcloudCuboidAnnotation):
track_id: int
is_keyframe: bool = False
[docs]class PointcloudSequenceCuboidFrame(PointcloudCuboidLabelAttributes):
annotations: List[PointcloudSequenceCuboidAnnotation]
timestamp: Optional[Union[str, int]] = None
format_version: Optional[FormatVersion] = None
[docs]class PointcloudSequenceCuboidLabelAttributes(BaseModel):
frames: List[PointcloudSequenceCuboidFrame]
format_version: Optional[FormatVersion] = None
# Point cloud sequence vector
[docs]class PointcloudSequenceVectorAnnotation(PointcloudVectorAnnotation):
track_id: int
is_keyframe: bool = False
[docs]class PointcloudSequenceVectorFrame(PointcloudVectorLabelAttributes):
annotations: List[PointcloudSequenceVectorAnnotation]
format_version: Optional[FormatVersion] = None
timestamp: Optional[Union[str, int]] = None
[docs]class PointcloudSequenceVectorLabelAttributes(BaseModel):
frames: List[PointcloudSequenceVectorFrame]
format_version: Optional[FormatVersion] = None
# Multi-sensor
[docs]class MultiSensorPointcloudSequenceCuboidLabelAttributes(BaseModel):
name: str
task_type: Literal[TaskType.POINTCLOUD_CUBOID_SEQUENCE]
attributes: PointcloudSequenceCuboidLabelAttributes
[docs]class MultiSensorImageSequenceVectorLabelAttributes(BaseModel):
name: str
task_type: Literal[TaskType.IMAGE_VECTOR_SEQUENCE]
attributes: ImageSequenceVectorLabelAttributes
[docs]class MultiSensorLabelAttributes(BaseModel):
sensors: List[
Union[
MultiSensorPointcloudSequenceCuboidLabelAttributes,
MultiSensorImageSequenceVectorLabelAttributes,
],
]
# Text
[docs]class TextAnnotation(BaseModel):
start: int
end: int
category_id: int
[docs]class TextLabelAttributes(BaseModel):
annotations: List[TextAnnotation]
format_version: Optional[FormatVersion] = None
# Most specific type first
# https://pydantic-docs.helpmanual.io/usage/types/#unions
LabelAttributes = Union[
ImageVectorLabelAttributes,
ImageSegmentationLabelAttributes,
ImageSequenceVectorLabelAttributes,
ImageSequenceSegmentationLabelAttributes,
PointcloudCuboidLabelAttributes,
PointcloudVectorLabelAttributes,
PointcloudSegmentationLabelAttributes,
PointcloudSequenceCuboidLabelAttributes,
PointcloudSequenceVectorLabelAttributes,
PointcloudSequenceSegmentationLabelAttributes,
MultiSensorLabelAttributes,
TextLabelAttributes,
]
[docs]class Label(BaseModel):
sample_uuid: str
label_type: TaskType
label_status: LabelStatus
labelset: str
attributes: Optional[LabelAttributes] = None
created_at: str
created_by: str
updated_at: str
score: Optional[float] = None
rating: Optional[float] = None
reviewed_at: Optional[str] = None
reviewed_by: Optional[str] = None
##########
# Sample #
##########
# Image
[docs]class ImageSampleAttributes(BaseModel):
image: URL
# Image sequence
[docs]class ImageFrame(ImageSampleAttributes):
name: Optional[str] = None
[docs]class ImageSequenceSampleAttributes(BaseModel):
frames: List[ImageFrame]
# Point cloud
[docs]class PCD(BaseModel):
url: str
signed_url: Optional[str] = None
type: PCDType
[docs]class EgoPose(BaseModel):
position: XYZ
heading: XYZW
[docs]class CameraIntrinsics(BaseModel):
intrinsic_matrix: List[List[float]]
[docs]class CameraExtrinsics(BaseModel):
translation: XYZ
rotation: XYZW
[docs]class CalibratedImage(URL):
row: Optional[int] = None
col: Optional[int] = None
intrinsics: Optional[CameraIntrinsics] = None
extrinsics: Optional[CameraExtrinsics] = None
distortion: Optional[Distortion] = None
camera_convention: CameraConvention = CameraConvention.OPEN_GL
[docs]class PointcloudSampleAttributes(BaseModel):
pcd: PCD
images: Optional[List[CalibratedImage]] = None
ego_pose: Optional[EgoPose] = None
default_z: Optional[float] = None
name: Optional[str] = None
timestamp: Optional[Union[str, int]] = None
# Point cloud sequence
[docs]class PointcloudSequenceSampleAttributes(BaseModel):
frames: List[PointcloudSampleAttributes]
# Multi-sensor
[docs]class MultiSensorPointcloudSequenceSampleAttributes(BaseModel):
name: str
task_type: Literal[TaskType.POINTCLOUD_CUBOID_SEQUENCE]
attributes: PointcloudSequenceSampleAttributes
[docs]class MultiSensorImageSequenceSampleAttributes(BaseModel):
name: str
task_type: Literal[TaskType.IMAGE_VECTOR_SEQUENCE]
attributes: ImageSequenceSampleAttributes
[docs]class MultiSensorSampleAttributes(BaseModel):
sensors: List[
Union[
MultiSensorPointcloudSequenceSampleAttributes,
MultiSensorImageSequenceSampleAttributes,
],
]
# Text
[docs]class TextSampleAttributes(BaseModel):
text: str
SampleAttributes = Union[
ImageSampleAttributes,
ImageSequenceSampleAttributes,
PointcloudSampleAttributes,
PointcloudSequenceSampleAttributes,
MultiSensorSampleAttributes,
TextSampleAttributes,
]
[docs]class Sample(BaseModel):
uuid: str
name: str
attributes: SampleAttributes
metadata: Dict[str, Any]
created_at: str
created_by: str
assigned_labeler: Optional[str] = None
assigned_reviewer: Optional[str] = None
comments: Optional[List[str]] = None
priority: float
has_embedding: Optional[bool] = None
label: Optional[Label] = None
issues: Optional[List[Issue]] = None
dataset_full_name: Optional[str] = None
########################
# Dataset and labelset #
########################
[docs]class User(BaseModel):
username: str
created_at: str
[docs]class Collaborator(BaseModel):
user: User
role: Role
[docs]class SelectTaskAttribute(BaseModel):
name: str
input_type: Literal[InputType.SELECT]
values: List[str]
default_value: Optional[str] = None
is_mandatory: Optional[bool] = None
[docs]class TextTaskAttribute(BaseModel):
name: str
input_type: Literal[InputType.TEXT] = None
default_value: Optional[str] = None
is_mandatory: Optional[bool] = None
[docs]class NumberTaskAttribute(BaseModel):
name: str
input_type: Literal[InputType.NUMBER]
default_value: Optional[float] = None
min: Optional[float] = None
max: Optional[float] = None
step: Optional[float] = None
is_mandatory: Optional[bool] = None
[docs] @field_validator("min", "max", "step", mode="before")
@classmethod
def empty_str_to_none(cls, v):
# min, max and step are empty strings when not filled in
if isinstance(v, str) and v.strip() == "":
return None
return v
[docs]class CheckboxTaskAttribute(BaseModel):
name: str
input_type: Literal[InputType.CHECKBOX]
default_value: Optional[bool] = None
TaskAttribute = Union[
SelectTaskAttribute,
TextTaskAttribute,
NumberTaskAttribute,
CheckboxTaskAttribute,
]
[docs]class TaskAttributeCategory(BaseModel):
name: str
id: int
color: Optional[Union[RGB, RGBA]] = None
has_instances: Optional[bool] = None
attributes: Optional[List[TaskAttribute]] = None
dimensions: Optional[XYZ] = None
model_config = ConfigDict(extra="allow")
[docs]class TaskAttributes(BaseModel):
format_version: Optional[FormatVersion] = None
categories: Optional[List[TaskAttributeCategory]] = None
image_attributes: Optional[List[TaskAttribute]] = None
model_config = ConfigDict(extra="allow")
[docs]class Owner(BaseModel):
username: str
created_at: str
email: Optional[str] = None
# class Statistics(BaseModel): # deprecated
# prelabeled_count: int
# labeled_count: int
# reviewed_count: int
# rejected_count: int
# skipped_count: int
# samples_count: int
[docs]class Labelset(BaseModel):
name: str
description: str
# uuid: Optional[str]
# readme: Optional[str]
# task_type: Optional[TaskType]
# attributes: Optional[Union[str, TaskAttributes]]
is_groundtruth: Optional[bool] = None
# statistics: Optional[Statistics]
created_at: Optional[str] = None
# stats: Optional[Dict[str, Any]]
[docs]class LabelStats(BaseModel):
TOTAL: Optional[int] = None
LABELED: Optional[int] = None
UNLABELED: Optional[int] = None
PRELABELED: Optional[int] = None
REVIEWED: Optional[int] = None
SKIPPED: Optional[int] = None
[docs]class Dataset(BaseModel):
name: str
full_name: str
cloned_from: Optional[str] = None
description: str
# data_type: Literal["IMAGE"]
category: Category
public: bool
owner: Owner
created_at: str
enable_ratings: bool
enable_skip_labeling: bool
enable_skip_reviewing: bool
enable_save_button: bool
enable_label_status_verified: bool
enable_same_dimensions_track_constraint: bool
enable_interpolation: bool
task_type: TaskType
# task_readme: str
label_stats: LabelStats
samples_count: Optional[Union[str, int]] = None
collaborators_count: Optional[int] = None
task_attributes: Optional[TaskAttributes] = None
labelsets: Optional[List[Labelset]] = None
role: Optional[str] = None
readme: Optional[str] = None
metadata: Dict[str, Any]
noncollaborator_can_label: Optional[bool] = None
noncollaborator_can_review: Optional[bool] = None
insights_urls: Optional[Dict[str, str]] = None
# tasks: Optional[List[Dict[str, Any]]]
embeddings_enabled: Optional[bool] = None
@field_validator("category")
@classmethod
def check_category(cls, category: str) -> str:
if category not in Category and "custom-" not in category:
raise ValidationError(
f"The category should be one of {Category}, but is {category}."
)
return category
####################
# Segments dataset #
####################
class SegmentsDatasetCategory(BaseModel):
id: int
name: str
color: Optional[Union[RGB, RGBA]] = None
attributes: Optional[List[Any]] = None
model_config = ConfigDict(populate_by_name=True, extra="allow")