helpy4/immich
1
0
Fork 0
mirror of https://github.com/immich-app/immich synced 2025-10-17 18:19:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

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
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
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()