# backward compatibility <=3.11: typing import TypedDict
from typing_extensions import TypedDict

from ai_parade._toolkit.structural_mapping import MappingDefinition

Output = MappingDefinition


class Classification(TypedDict, total=False):
	label: list[str]

	class_ImageNet1k: int
	class_ImageNet1k_id: str
	class_COCO: int

	class_ImageNet1k_logit: list[float]
	class_COCO_logit: list[float]

	class_ImageNet1k_probability: list[float]
	class_COCO_probability: list[float]


class Detection(Classification, total=False):
	bbox_xmin_absolute: float
	bbox_xmax_absolute: float
	bbox_ymin_absolute: float
	bbox_ymax_absolute: float

	bbox_xcenter_absolute: float
	bbox_ycenter_absolute: float
	bbox_width_absolute: float
	bbox_height_absolute: float

	bbox_xmin_normalized: float
	bbox_xmax_normalized: float
	bbox_ymin_normalized: float
	bbox_ymax_normalized: float

	bbox_xcenter_normalized: float
	bbox_ycenter_normalized: float
	bbox_width_normalized: float
	bbox_height_normalized: float


class ModelOutput(Classification, total=False):
	"""
	Represents the output of the model.

	Currently only vision models are supported.
	"""

	detections: list[Detection]
