1
0
mirror of https://github.com/immich-app/immich.git synced 2024-12-27 10:58:13 +02:00
immich/machine-learning/app/schemas.py
Mert a2f5674bbb
refactor(ml): modularization and styling (#2835)
* basic refactor and styling

* removed batching

* module entrypoint

* removed unused imports

* model superclass,  model cache now in app state

* fixed cache dir and enforced abstract method

---------

Co-authored-by: Alex Tran <alex.tran1502@gmail.com>
2023-06-24 22:18:09 -05:00

67 lines
1.1 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"
CLIP_VISION = "clip-vision"
CLIP_TEXT = "clip-text"
FACIAL_RECOGNITION = "facial-recognition"