1
0
Fork 0
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 ()

* env variables for tags, faces and eager startup

* chore(server,ml): remove object detection job and endpoint ()

* 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:
Mert 2023-06-02 22:42:47 -04:00 committed by GitHub
parent c5234731d6
commit b8de668f5f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -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))