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

feat(ml) backend takes image over HTTP (#2783)

Browse source 
* using pydantic BaseSetting

* ML API takes image file as input

* keeping image in memory

* reducing duplicate code

* using bytes instead of UploadFile & other small code improvements

* removed form-multipart, using HTTP body

* format code

---------

Co-authored-by: Alex Tran <alex.tran1502@gmail.com>
This commit is contained in:
Zeeshan Khan 2023-06-17 22:49:19 -05:00 • committed by GitHub
parent 3e804f16df
commit 34201be74c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 116 additions and 80 deletions
Download patch file Download diff file Expand all files Collapse all files

20
machine-learning/app/models.py
View file

@ -1,14 +1,15 @@
import torch
from insightface.app import FaceAnalysis
from pathlib import Path
import os
from transformers import pipeline, Pipeline
from sentence_transformers import SentenceTransformer
from typing import Any
from typing import Any, BinaryIO
import cv2 as cv
import numpy as np
from PIL import Image
from config import settings
cache_folder = os.getenv("MACHINE_LEARNING_CACHE_FOLDER", "/cache")
device = "cuda" if torch.cuda.is_available() else "cpu"
@ -49,9 +50,9 @@ def get_model(model_name: str, model_type: str, **model_kwargs):
def run_classification(
model: Pipeline, image_path: str, min_score: float | None = None
model: Pipeline, image: Image, min_score: float | None = None
):
predictions: list[dict[str, Any]] = model(image_path) # type: ignore
predictions: list[dict[str, Any]] = model(image) # type: ignore
result = {
tag
for pred in predictions
@ -63,9 +64,10 @@ def run_classification(
def run_facial_recognition(
model: FaceAnalysis, image_path: str
model: FaceAnalysis, image: bytes
) -> list[dict[str, Any]]:
img = cv.imread(image_path)
file_bytes = np.frombuffer(image, dtype=np.uint8)
img = cv.imdecode(file_bytes, cv.IMREAD_COLOR)
height, width, _ = img.shape
results = []
faces = model.get(img)
@ -101,7 +103,7 @@ def _load_facial_recognition(
if isinstance(cache_dir, Path):
cache_dir = cache_dir.as_posix()
if min_face_score is None:
min_face_score = float(os.getenv("MACHINE_LEARNING_MIN_FACE_SCORE", 0.7))
min_face_score = settings.min_face_score
model = FaceAnalysis(
name=model_name,
@ -114,4 +116,4 @@ def _load_facial_recognition(
def _get_cache_dir(model_name: str, model_type: str) -> Path:
return Path(cache_folder, device, model_type, model_name)
return Path(settings.cache_folder, device, model_type, model_name)