mirror of
https://github.com/immich-app/immich.git
synced 2025-01-27 22:22:45 +01:00
feat(ml): env variables for tags, faces and eager startup (#2626)
* env variables for tags, faces and eager startup * chore(server,ml): remove object detection job and endpoint (#2627) * removed object detection job * removed object detection endpoint * env variables for tags, faces and eager startup * download without caching models if not eager * simplified `get_cached_model` * re-added env for clip text model
This commit is contained in:
parent
c5234731d6
commit
b8de668f5f
1 changed files with 45 additions and 24 deletions
|
@ -28,6 +28,12 @@ facial_recognition_model = os.getenv(
|
||||||
"MACHINE_LEARNING_FACIAL_RECOGNITION_MODEL", "buffalo_l"
|
"MACHINE_LEARNING_FACIAL_RECOGNITION_MODEL", "buffalo_l"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
min_face_score = float(os.getenv("MACHINE_LEARNING_MIN_FACE_SCORE", 0.7))
|
||||||
|
min_tag_score = float(os.getenv("MACHINE_LEARNING_MIN_TAG_SCORE", 0.9))
|
||||||
|
eager_startup = (
|
||||||
|
os.getenv("MACHINE_LEARNING_EAGER_STARTUP", "true") == "true"
|
||||||
|
) # loads all models at startup
|
||||||
|
|
||||||
cache_folder = os.getenv("MACHINE_LEARNING_CACHE_FOLDER", "/cache")
|
cache_folder = os.getenv("MACHINE_LEARNING_CACHE_FOLDER", "/cache")
|
||||||
|
|
||||||
_model_cache = {}
|
_model_cache = {}
|
||||||
|
@ -37,11 +43,19 @@ app = FastAPI()
|
||||||
|
|
||||||
@app.on_event("startup")
|
@app.on_event("startup")
|
||||||
async def startup_event():
|
async def startup_event():
|
||||||
|
models = [
|
||||||
|
(classification_model, "image-classification"),
|
||||||
|
(clip_image_model, "clip"),
|
||||||
|
(clip_text_model, "clip"),
|
||||||
|
(facial_recognition_model, "facial-recognition"),
|
||||||
|
]
|
||||||
|
|
||||||
# Get all models
|
# Get all models
|
||||||
_get_model(classification_model, "image-classification")
|
for model_name, model_type in models:
|
||||||
_get_model(clip_image_model)
|
if eager_startup:
|
||||||
_get_model(clip_text_model)
|
get_cached_model(model_name, model_type)
|
||||||
_get_model(facial_recognition_model, "facial-recognition")
|
else:
|
||||||
|
_get_model(model_name, model_type)
|
||||||
|
|
||||||
|
|
||||||
@app.get("/")
|
@app.get("/")
|
||||||
|
@ -53,30 +67,31 @@ async def root():
|
||||||
def ping():
|
def ping():
|
||||||
return "pong"
|
return "pong"
|
||||||
|
|
||||||
|
|
||||||
@app.post("/image-classifier/tag-image", status_code=200)
|
@app.post("/image-classifier/tag-image", status_code=200)
|
||||||
def image_classification(payload: MlRequestBody):
|
def image_classification(payload: MlRequestBody):
|
||||||
model = _get_model(classification_model, "image-classification")
|
model = get_cached_model(classification_model, "image-classification")
|
||||||
assetPath = payload.thumbnailPath
|
assetPath = payload.thumbnailPath
|
||||||
return run_engine(model, assetPath)
|
return run_engine(model, assetPath)
|
||||||
|
|
||||||
|
|
||||||
@app.post("/sentence-transformer/encode-image", status_code=200)
|
@app.post("/sentence-transformer/encode-image", status_code=200)
|
||||||
def clip_encode_image(payload: MlRequestBody):
|
def clip_encode_image(payload: MlRequestBody):
|
||||||
model = _get_model(clip_image_model)
|
model = get_cached_model(clip_image_model, "clip")
|
||||||
assetPath = payload.thumbnailPath
|
assetPath = payload.thumbnailPath
|
||||||
return model.encode(Image.open(assetPath)).tolist()
|
return model.encode(Image.open(assetPath)).tolist()
|
||||||
|
|
||||||
|
|
||||||
@app.post("/sentence-transformer/encode-text", status_code=200)
|
@app.post("/sentence-transformer/encode-text", status_code=200)
|
||||||
def clip_encode_text(payload: ClipRequestBody):
|
def clip_encode_text(payload: ClipRequestBody):
|
||||||
model = _get_model(clip_text_model)
|
model = get_cached_model(clip_text_model, "clip")
|
||||||
text = payload.text
|
text = payload.text
|
||||||
return model.encode(text).tolist()
|
return model.encode(text).tolist()
|
||||||
|
|
||||||
|
|
||||||
@app.post("/facial-recognition/detect-faces", status_code=200)
|
@app.post("/facial-recognition/detect-faces", status_code=200)
|
||||||
def facial_recognition(payload: MlRequestBody):
|
def facial_recognition(payload: MlRequestBody):
|
||||||
model = _get_model(facial_recognition_model, "facial-recognition")
|
model = get_cached_model(facial_recognition_model, "facial-recognition")
|
||||||
assetPath = payload.thumbnailPath
|
assetPath = payload.thumbnailPath
|
||||||
img = cv.imread(assetPath)
|
img = cv.imread(assetPath)
|
||||||
height, width, _ = img.shape
|
height, width, _ = img.shape
|
||||||
|
@ -84,7 +99,7 @@ def facial_recognition(payload: MlRequestBody):
|
||||||
faces = model.get(img)
|
faces = model.get(img)
|
||||||
|
|
||||||
for face in faces:
|
for face in faces:
|
||||||
if face.det_score < 0.7:
|
if face.det_score < min_face_score:
|
||||||
continue
|
continue
|
||||||
x1, y1, x2, y2 = face.bbox
|
x1, y1, x2, y2 = face.bbox
|
||||||
|
|
||||||
|
@ -111,7 +126,7 @@ def run_engine(engine, path):
|
||||||
|
|
||||||
for index, pred in enumerate(predictions):
|
for index, pred in enumerate(predictions):
|
||||||
tags = pred["label"].split(", ")
|
tags = pred["label"].split(", ")
|
||||||
if pred["score"] > 0.9:
|
if pred["score"] > min_tag_score:
|
||||||
result = [*result, *tags]
|
result = [*result, *tags]
|
||||||
|
|
||||||
if len(result) > 1:
|
if len(result) > 1:
|
||||||
|
@ -120,26 +135,32 @@ def run_engine(engine, path):
|
||||||
return result
|
return result
|
||||||
|
|
||||||
|
|
||||||
def _get_model(model, task=None):
|
def get_cached_model(model, task):
|
||||||
global _model_cache
|
global _model_cache
|
||||||
key = "|".join([model, str(task)])
|
key = "|".join([model, str(task)])
|
||||||
if key not in _model_cache:
|
if key not in _model_cache:
|
||||||
if task:
|
model = _get_model(model, task)
|
||||||
if task == "facial-recognition":
|
_model_cache[key] = model
|
||||||
face_model = FaceAnalysis(
|
|
||||||
name=model,
|
|
||||||
root=cache_folder,
|
|
||||||
allowed_modules=["detection", "recognition"],
|
|
||||||
)
|
|
||||||
face_model.prepare(ctx_id=0, det_size=(640, 640))
|
|
||||||
_model_cache[key] = face_model
|
|
||||||
else:
|
|
||||||
_model_cache[key] = pipeline(model=model, task=task)
|
|
||||||
else:
|
|
||||||
_model_cache[key] = SentenceTransformer(model, cache_folder=cache_folder)
|
|
||||||
return _model_cache[key]
|
return _model_cache[key]
|
||||||
|
|
||||||
|
|
||||||
|
def _get_model(model, task):
|
||||||
|
match task:
|
||||||
|
case "facial-recognition":
|
||||||
|
model = FaceAnalysis(
|
||||||
|
name=model,
|
||||||
|
root=cache_folder,
|
||||||
|
allowed_modules=["detection", "recognition"],
|
||||||
|
)
|
||||||
|
model.prepare(ctx_id=0, det_size=(640, 640))
|
||||||
|
case "clip":
|
||||||
|
model = SentenceTransformer(model, cache_folder=cache_folder)
|
||||||
|
case _:
|
||||||
|
model = pipeline(model=model, task=task)
|
||||||
|
return model
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
host = os.getenv("MACHINE_LEARNING_HOST", "0.0.0.0")
|
host = os.getenv("MACHINE_LEARNING_HOST", "0.0.0.0")
|
||||||
port = int(os.getenv("MACHINE_LEARNING_PORT", 3003))
|
port = int(os.getenv("MACHINE_LEARNING_PORT", 3003))
|
||||||
|
|
Loading…
Reference in a new issue