schemas.py 984 B

12345678910111213141516171819202122232425262728293031323334353637383940414243444546
  1. from enum import StrEnum
  2. from typing import Any, Protocol, TypeAlias, TypedDict, TypeGuard
  3. import numpy as np
  4. from pydantic import BaseModel
  5. ndarray_f32: TypeAlias = np.ndarray[int, np.dtype[np.float32]]
  6. ndarray_i64: TypeAlias = np.ndarray[int, np.dtype[np.int64]]
  7. ndarray_i32: TypeAlias = np.ndarray[int, np.dtype[np.int32]]
  8. class TextResponse(BaseModel):
  9. __root__: str
  10. class MessageResponse(BaseModel):
  11. message: str
  12. class BoundingBox(TypedDict):
  13. x1: int
  14. y1: int
  15. x2: int
  16. y2: int
  17. class ModelType(StrEnum):
  18. IMAGE_CLASSIFICATION = "image-classification"
  19. CLIP = "clip"
  20. FACIAL_RECOGNITION = "facial-recognition"
  21. class HasProfiling(Protocol):
  22. profiling: dict[str, float]
  23. class Face(TypedDict):
  24. boundingBox: BoundingBox
  25. embedding: ndarray_f32
  26. imageWidth: int
  27. imageHeight: int
  28. score: float
  29. def has_profiling(obj: Any) -> TypeGuard[HasProfiling]:
  30. return hasattr(obj, "profiling") and type(obj.profiling) == dict