immich/machine-learning/app/schemas.py
Mert df1e8679d9
chore(ml): added testing and github workflow (#2969)
* added testing

* github action for python, made mypy happy

* formatted with black

* minor fixes and styling

* test model cache

* cache test dependencies

* narrowed model cache tests

* moved endpoint tests to their own class

* cleaned up fixtures

* formatting

* removed unused dep
2023-06-27 18:21:33 -05:00

62 lines
1.0 KiB
Python

from enum import Enum
from pydantic import BaseModel
def to_lower_camel(string: str) -> str:
tokens = [token.capitalize() if i > 0 else token for i, token in enumerate(string.split("_"))]
return "".join(tokens)
class TextModelRequest(BaseModel):
text: str
class TextResponse(BaseModel):
__root__: str
class MessageResponse(BaseModel):
message: str
class TagResponse(BaseModel):
__root__: list[str]
class Embedding(BaseModel):
__root__: list[float]
class EmbeddingResponse(BaseModel):
__root__: Embedding
class BoundingBox(BaseModel):
x1: int
y1: int
x2: int
y2: int
class Face(BaseModel):
image_width: int
image_height: int
bounding_box: BoundingBox
score: float
embedding: Embedding
class Config:
alias_generator = to_lower_camel
allow_population_by_field_name = True
class FaceResponse(BaseModel):
__root__: list[Face]
class ModelType(Enum):
IMAGE_CLASSIFICATION = "image-classification"
CLIP = "clip"
FACIAL_RECOGNITION = "facial-recognition"