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,
# Whether to populate models with the value property of enums, rather than the raw enum
use_enum_values=True,
)
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"
VERIFIED = "VERIFIED"
UNLABELED = "UNLABELED"
# keep in sync with LabelStatus
[docs]class LabelStats(BaseModel):
REVIEWED: Optional[int] = None
REVIEWING_IN_PROGRESS: Optional[int] = None
LABELED: Optional[int] = None
LABELING_IN_PROGRESS: Optional[int] = None
REJECTED: Optional[int] = None
PRELABELED: Optional[int] = None
SKIPPED: Optional[int] = None
VERIFIED: Optional[int] = None
UNLABELED: Optional[int] = None
# extra
TOTAL: Optional[int] = None
[docs]class TaskType(str, Enum):
SEGMENTATION_BITMAP = "segmentation-bitmap"
SEGMENTATION_BITMAP_HIGHRES = "segmentation-bitmap-highres"
IMAGE_SEGMENTATION_SEQUENCE = "image-segmentation-sequence"
BBOXES = "bboxes"
VECTOR = "vector"
IMAGE_VECTOR_SEQUENCE = "image-vector-sequence"
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" # combination of pointcloud-cuboid-sequence and image-vector-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"
VECTOR3 = "vector3"
QUATERNION = "quaternion"
POINTS = "points"
[docs]class CameraDistortionModel(str, Enum):
FISH_EYE = "fisheye"
BROWN_CONRADY = "brown-conrady"
[docs]class PCDType(str, Enum):
PCD = "pcd"
BINARY_XYZI = "binary-xyzi"
KITTI = "kitti"
BINARY_XYZIR = "binary-xyzir"
NUSCENES = "nuscenes"
PLY = "ply"
LAS = "las"
[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"
class Subscription(str, Enum):
FREE = "FREE"
STANDARD = "STANDARD"
ENTERPRISE = "ENTERPRISE"
ACADEMIC = "ACADEMIC"
TRIAL = "TRIAL"
RGB = Tuple[int, int, int]
RGBA = Tuple[int, int, int, int]
FormatVersion = Union[float, str]
ObjectAttributes = Dict[str, Optional[Union[str, bool, int, float]]]
ImageAttributes = Dict[str, Optional[Union[str, bool, int, float]]]
Timestamp = Union[str, int, float]
###########
# Release #
###########
[docs]class URL(BaseModel):
# TODO Remove optional (the backend does not return an URL when adding a release)
url: Optional[str] = None
[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[Timestamp] = None
format_version: Optional[FormatVersion] = None
class ImageSequenceSegmentationLabelAttributes(BaseModel):
frames: List[ImageSequenceSegmentationFrame]
format_version: Optional[FormatVersion] = None
# Image sequence vector
[docs]class ImageSequenceVectorAnnotation(ImageVectorAnnotation):
track_id: int
is_keyframe: bool = False
[docs]class ImageVectorFrame(ImageVectorLabelAttributes):
annotations: List[ImageSequenceVectorAnnotation]
timestamp: Optional[Timestamp] = 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[Timestamp] = 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[Timestamp] = 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[Timestamp] = 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]
# TODO remove list and replace with `Optional[PointcloudSequenceCuboidLabelAttributes] = None`
attributes: Union[PointcloudSequenceCuboidLabelAttributes, List]
[docs]class MultiSensorImageSequenceVectorLabelAttributes(BaseModel):
name: str
task_type: Literal[TaskType.IMAGE_VECTOR_SEQUENCE]
# TODO remove list and replace with `Optional[ImageSequenceVectorLabelAttributes] = None`
attributes: Union[ImageSequenceVectorLabelAttributes, List]
[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
# 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
class LabelSummary(BaseModel):
score: Optional[float] = None
label_status: LabelStatus
updated_at: Optional[str] = None
created_by: 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
timestamp: Optional[Timestamp] = 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
class Bounds(BaseModel):
min_x: Optional[float] = None
max_x: Optional[float] = None
min_y: Optional[float] = None
max_y: Optional[float] = None
min_z: Optional[float] = None
max_z: Optional[float] = None
[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
name: Optional[str] = None
rotation: Optional[float] = None
[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[Timestamp] = None
bounds: Optional[Bounds] = 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
label: Optional[Union[Label, LabelSummary]] = None
issues: Optional[List[Issue]] = None
dataset_full_name: Optional[str] = None
########################
# Dataset and labelset #
########################
# https://docs.pydantic.dev/latest/concepts/postponed_annotations/#self-referencing-or-recursive-models
[docs]class User(BaseModel):
username: str
created_at: str
is_organization: bool
email: Optional[str] = None
webhooks_enabled: Optional[bool] = None
private_upload_count: Optional[int] = None
public_upload_count: Optional[int] = None
subscription: Optional[Subscription] = None
is_trial_expired: Optional[bool] = None
organizations: Optional[List[User]] = None
organization_created_by: Optional[str] = None
organization_role: Optional[Role] = None
members: Optional[List[User]] = None
insights_urls: Optional[Dict[str, str]] = None
[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
is_track_level: Optional[bool] = None
[docs]class TextTaskAttribute(BaseModel):
name: str
input_type: Literal[InputType.TEXT]
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
is_track_level: Optional[bool] = None
class Vector3TaskAttribute(BaseModel):
name: str
input_type: Literal[InputType.VECTOR3]
is_mandatory: Optional[bool] = None
is_track_level: Optional[bool] = None
class QuaternionTaskAttribute(BaseModel):
name: str
input_type: Literal[InputType.QUATERNION]
is_mandatory: Optional[bool] = None
is_track_level: Optional[bool] = None
class PointsTaskAttribute(BaseModel):
name: str
input_type: Literal[InputType.POINTS]
is_mandatory: Optional[bool] = None
is_track_level: Optional[bool] = None
TaskAttribute = Union[
SelectTaskAttribute,
TextTaskAttribute,
NumberTaskAttribute,
CheckboxTaskAttribute,
Vector3TaskAttribute,
QuaternionTaskAttribute,
PointsTaskAttribute,
]
[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] = None
# readme: Optional[str] = None
# task_type: Optional[TaskType] = None
# attributes: Optional[Union[str, TaskAttributes]] = None
is_groundtruth: Optional[bool] = None
# statistics: Optional[Statistics] = None
created_at: Optional[str] = None
# stats: Optional[Dict[str, Any]] = 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
use_timestamps_for_interpolation: bool
task_type: TaskType
# task_readme: str
label_stats: LabelStats
labeling_inactivity_timeout_seconds: Optional[int] = None
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]]] = 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")