from pydantic import BaseModel
from typing import Literal, Set, Optional


class ImageDataContract(BaseModel):
    data_type: Literal["image"] = "image"
    modality: Literal["CT", "PET", "MR"]
    dim: Literal["2D", "3D", "2_5D"]
    geometry_id: Optional[str] = None
    has_mask: bool = False


class MaskDataContract(BaseModel):
    data_type: Literal["mask"] = "mask"
    modality: Literal["CT", "PET", "MR"]
    dim: Literal["2D", "3D", "2_5D"]
    geometry_id: Optional[str] = None


class RadiomicsFeatureTableContract(BaseModel):
    data_type: Literal["table"] = "table"
