ashipov/immich
1
0
Fork 1
forked from dark_thunder/immich
Code Pull Requests Activity

feat(ml)!: switch image classification and CLIP models to ONNX (#3809)

Browse Source
This commit is contained in:
Mert
2023-08-25 00:28:51 -04:00
committed by GitHub
parent 8211afb726
commit 165b91b068
14 changed files with 1617 additions and 507 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
machine-learning/app/config.py
View File

@ -1,3 +1,4 @@
import os
from pathlib import Path
from pydantic import BaseSettings
@ -8,8 +9,8 @@ from .schemas import ModelType
class Settings(BaseSettings):
cache_folder: str = "/cache"
classification_model: str = "microsoft/resnet-50"
clip_image_model: str = "clip-ViT-B-32"
clip_text_model: str = "clip-ViT-B-32"
clip_image_model: str = "ViT-B-32::openai"
clip_text_model: str = "ViT-B-32::openai"
facial_recognition_model: str = "buffalo_l"
min_tag_score: float = 0.9
eager_startup: bool = False
@ -19,14 +20,20 @@ class Settings(BaseSettings):
workers: int = 1
min_face_score: float = 0.7
test_full: bool = False
request_threads: int = os.cpu_count() or 4
model_inter_op_threads: int = 1
model_intra_op_threads: int = 2
class Config:
env_prefix = "MACHINE_LEARNING_"
case_sensitive = False
_clean_name = str.maketrans(":\\/", "___", ".")
def get_cache_dir(model_name: str, model_type: ModelType) -> Path:
return Path(settings.cache_folder, model_type.value, model_name)
return Path(settings.cache_folder) / model_type.value / model_name.translate(_clean_name)
settings = Settings()