refactor: migrate person repository to kysely (#15242)

* refactor: migrate person repository to kysely

* `asVector` begone

* linting

* fix metadata faces

* update test

---------

Co-authored-by: Alex <[email protected]>
Co-authored-by: mertalev <[email protected]>

e2e/src/api/specs/person.e2e-spec.ts

@@ -200,7 +200,7 @@ describe('/people', () => {
       expect(body).toMatchObject({
         id: expect.any(String),
         name: 'New Person',
-        birthDate: '1990-01-01',
+        birthDate: '1990-01-01T00:00:00.000Z',
       });
     });
   });
@@ -244,7 +244,7 @@ describe('/people', () => {
         .set('Authorization', `Bearer ${admin.accessToken}`)
         .send({ birthDate: '1990-01-01' });
       expect(status).toBe(200);
-      expect(body).toMatchObject({ birthDate: '1990-01-01' });
+      expect(body).toMatchObject({ birthDate: '1990-01-01T00:00:00.000Z' });
     });
 
     it('should clear a date of birth', async () => {

machine-learning/app/models/clip/textual.py

@@ -10,17 +10,17 @@
 
 from app.config import log
 from app.models.base import InferenceModel
-from app.models.transforms import clean_text
+from app.models.transforms import clean_text, serialize_np_array
 from app.schemas import ModelSession, ModelTask, ModelType
 
 
 class BaseCLIPTextualEncoder(InferenceModel):
     depends = []
     identity = (ModelType.TEXTUAL, ModelTask.SEARCH)
 
-    def _predict(self, inputs: str, **kwargs: Any) -> NDArray[np.float32]:
+    def _predict(self, inputs: str, **kwargs: Any) -> str:
         res: NDArray[np.float32] = self.session.run(None, self.tokenize(inputs))[0][0]
-        return res
+        return serialize_np_array(res)
 
     def _load(self) -> ModelSession:
         session = super()._load()

machine-learning/app/models/clip/visual.py

@@ -10,18 +10,26 @@
 
 from app.config import log
 from app.models.base import InferenceModel
-from app.models.transforms import crop_pil, decode_pil, get_pil_resampling, normalize, resize_pil, to_numpy
+from app.models.transforms import (
+    crop_pil,
+    decode_pil,
+    get_pil_resampling,
+    normalize,
+    resize_pil,
+    serialize_np_array,
+    to_numpy,
+)
 from app.schemas import ModelSession, ModelTask, ModelType
 
 
 class BaseCLIPVisualEncoder(InferenceModel):
     depends = []
     identity = (ModelType.VISUAL, ModelTask.SEARCH)
 
-    def _predict(self, inputs: Image.Image | bytes, **kwargs: Any) -> NDArray[np.float32]:
+    def _predict(self, inputs: Image.Image | bytes, **kwargs: Any) -> str:
         image = decode_pil(inputs)
         res: NDArray[np.float32] = self.session.run(None, self.transform(image))[0][0]
-        return res
+        return serialize_np_array(res)
 
     @abstractmethod
     def transform(self, image: Image.Image) -> dict[str, NDArray[np.float32]]:

machine-learning/app/models/facial_recognition/recognition.py

@@ -12,7 +12,7 @@
 
 from app.config import log, settings
 from app.models.base import InferenceModel
-from app.models.transforms import decode_cv2
+from app.models.transforms import decode_cv2, serialize_np_array
 from app.schemas import FaceDetectionOutput, FacialRecognitionOutput, ModelFormat, ModelSession, ModelTask, ModelType
 
 
@@ -61,7 +61,7 @@ def postprocess(self, faces: FaceDetectionOutput, embeddings: NDArray[np.float32
         return [
             {
                 "boundingBox": {"x1": x1, "y1": y1, "x2": x2, "y2": y2},
-                "embedding": embedding,
+                "embedding": serialize_np_array(embedding),
                 "score": score,
             }
             for (x1, y1, x2, y2), embedding, score in zip(faces["boxes"], embeddings, faces["scores"])

machine-learning/app/models/transforms.py

@@ -4,6 +4,7 @@
 
 import cv2
 import numpy as np
+import orjson
 from numpy.typing import NDArray
 from PIL import Image
 
@@ -69,3 +70,9 @@ def clean_text(text: str, canonicalize: bool = False) -> str:
     if canonicalize:
         text = text.translate(_PUNCTUATION_TRANS).lower()
     return text
+
+
+# this allows the client to use the array as a string without deserializing only to serialize back to a string
+# TODO: use this in a less invasive way
+def serialize_np_array(arr: NDArray[np.float32]) -> str:
+    return orjson.dumps(arr, option=orjson.OPT_SERIALIZE_NUMPY).decode()

machine-learning/app/schemas.py

@@ -79,7 +79,7 @@ class FaceDetectionOutput(TypedDict):
 
 class DetectedFace(TypedDict):
     boundingBox: BoundingBox
-    embedding: npt.NDArray[np.float32]
+    embedding: str
     score: float
 
 

machine-learning/app/test_main.py

@@ -10,6 +10,7 @@
 import cv2
 import numpy as np
 import onnxruntime as ort
+import orjson
 import pytest
 from fastapi import HTTPException
 from fastapi.testclient import TestClient
@@ -346,11 +347,11 @@ def test_basic_image(
         mocked.run.return_value = [[self.embedding]]
 
         clip_encoder = OpenClipVisualEncoder("ViT-B-32__openai", cache_dir="test_cache")
-        embedding = clip_encoder.predict(pil_image)
-
-        assert isinstance(embedding, np.ndarray)
-        assert embedding.shape[0] == clip_model_cfg["embed_dim"]
-        assert embedding.dtype == np.float32
+        embedding_str = clip_encoder.predict(pil_image)
+        assert isinstance(embedding_str, str)
+        embedding = orjson.loads(embedding_str)
+        assert isinstance(embedding, list)
+        assert len(embedding) == clip_model_cfg["embed_dim"]
         mocked.run.assert_called_once()
 
     def test_basic_text(
@@ -368,11 +369,11 @@ def test_basic_text(
         mocker.patch("app.models.clip.textual.Tokenizer.from_file", autospec=True)
 
         clip_encoder = OpenClipTextualEncoder("ViT-B-32__openai", cache_dir="test_cache")
-        embedding = clip_encoder.predict("test search query")
-
-        assert isinstance(embedding, np.ndarray)
-        assert embedding.shape[0] == clip_model_cfg["embed_dim"]
-        assert embedding.dtype == np.float32
+        embedding_str = clip_encoder.predict("test search query")
+        assert isinstance(embedding_str, str)
+        embedding = orjson.loads(embedding_str)
+        assert isinstance(embedding, list)
+        assert len(embedding) == clip_model_cfg["embed_dim"]
         mocked.run.assert_called_once()
 
     def test_openclip_tokenizer(
@@ -508,8 +509,11 @@ def test_recognition(self, cv_image: cv2.Mat, mocker: MockerFixture) -> None:
             assert isinstance(face.get("boundingBox"), dict)
             assert set(face["boundingBox"]) == {"x1", "y1", "x2", "y2"}
             assert all(isinstance(val, np.float32) for val in face["boundingBox"].values())
-            assert isinstance(face.get("embedding"), np.ndarray)
-            assert face["embedding"].shape[0] == 512
+            embedding_str = face.get("embedding")
+            assert isinstance(embedding_str, str)
+            embedding = orjson.loads(embedding_str)
+            assert isinstance(embedding, list)
+            assert len(embedding) == 512
             assert isinstance(face.get("score", None), np.float32)
 
         rec_model.get_feat.assert_called_once()
@@ -880,8 +884,10 @@ def test_clip_image_endpoint(
         actual = response.json()
         assert response.status_code == 200
         assert isinstance(actual, dict)
-        assert isinstance(actual.get("clip", None), list)
-        assert np.allclose(expected, actual["clip"])
+        embedding = actual.get("clip", None)
+        assert isinstance(embedding, str)
+        parsed_embedding = orjson.loads(embedding)
+        assert np.allclose(expected, parsed_embedding)
 
     def test_clip_text_endpoint(self, responses: dict[str, Any], deployed_app: TestClient) -> None:
         expected = responses["clip"]["text"]
@@ -901,8 +907,10 @@ def test_clip_text_endpoint(self, responses: dict[str, Any], deployed_app: TestC
         actual = response.json()
         assert response.status_code == 200
         assert isinstance(actual, dict)
-        assert isinstance(actual.get("clip", None), list)
-        assert np.allclose(expected, actual["clip"])
+        embedding = actual.get("clip", None)
+        assert isinstance(embedding, str)
+        parsed_embedding = orjson.loads(embedding)
+        assert np.allclose(expected, parsed_embedding)
 
     def test_face_endpoint(self, pil_image: Image.Image, responses: dict[str, Any], deployed_app: TestClient) -> None:
         byte_image = BytesIO()
@@ -933,5 +941,8 @@ def test_face_endpoint(self, pil_image: Image.Image, responses: dict[str, Any],
 
         for expected_face, actual_face in zip(responses["facial-recognition"], actual["facial-recognition"]):
             assert expected_face["boundingBox"] == actual_face["boundingBox"]
-            assert np.allclose(expected_face["embedding"], actual_face["embedding"])
+            embedding = actual_face.get("embedding", None)
+            assert isinstance(embedding, str)
+            parsed_embedding = orjson.loads(embedding)
+            assert np.allclose(expected_face["embedding"], parsed_embedding)
             assert np.allclose(expected_face["score"], actual_face["score"])

server/src/decorators.ts

@@ -100,6 +100,7 @@ export const DummyValue = {
   DATE: new Date(),
   TIME_BUCKET: '2024-01-01T00:00:00.000Z',
   BOOLEAN: true,
+  VECTOR: '[1, 2, 3]',
 };
 
 export const GENERATE_SQL_KEY = 'generate-sql-key';

server/src/entities/face-search.entity.ts

@@ -11,10 +11,6 @@ export class FaceSearchEntity {
   faceId!: string;
 
   @Index('face_index', { synchronize: false })
-  @Column({
-    type: 'float4',
-    array: true,
-    transformer: { from: JSON.parse, to: (v) => `[${v}]` },
-  })
-  embedding!: number[];
+  @Column({ type: 'float4', array: true })
+  embedding!: string;
 }

server/src/entities/smart-search.entity.ts

@@ -11,6 +11,6 @@ export class SmartSearchEntity {
   assetId!: string;
 
   @Index('clip_index', { synchronize: false })
-  @Column({ type: 'float4', array: true, transformer: { from: JSON.parse, to: (v) => v } })
-  embedding!: number[];
+  @Column({ type: 'float4', array: true })
+  embedding!: string;
 }

server/src/interfaces/machine-learning.interface.ts

@@ -28,10 +28,10 @@ export type FaceDetectionOptions = ModelOptions & { minScore: number };
 
 type VisualResponse = { imageHeight: number; imageWidth: number };
 export type ClipVisualRequest = { [ModelTask.SEARCH]: { [ModelType.VISUAL]: ModelOptions } };
-export type ClipVisualResponse = { [ModelTask.SEARCH]: number[] } & VisualResponse;
+export type ClipVisualResponse = { [ModelTask.SEARCH]: string } & VisualResponse;
 
 export type ClipTextualRequest = { [ModelTask.SEARCH]: { [ModelType.TEXTUAL]: ModelOptions } };
-export type ClipTextualResponse = { [ModelTask.SEARCH]: number[] };
+export type ClipTextualResponse = { [ModelTask.SEARCH]: string };
 
 export type FacialRecognitionRequest = {
   [ModelTask.FACIAL_RECOGNITION]: {
@@ -42,7 +42,7 @@ export type FacialRecognitionRequest = {
 
 export interface Face {
   boundingBox: BoundingBox;
-  embedding: number[];
+  embedding: string;
   score: number;
 }
 
@@ -51,7 +51,7 @@ export type DetectedFaces = { faces: Face[] } & VisualResponse;
 export type MachineLearningRequest = ClipVisualRequest | ClipTextualRequest | FacialRecognitionRequest;
 
 export interface IMachineLearningRepository {
-  encodeImage(urls: string[], imagePath: string, config: ModelOptions): Promise<number[]>;
-  encodeText(urls: string[], text: string, config: ModelOptions): Promise<number[]>;
+  encodeImage(urls: string[], imagePath: string, config: ModelOptions): Promise<string>;
+  encodeText(urls: string[], text: string, config: ModelOptions): Promise<string>;
   detectFaces(urls: string[], imagePath: string, config: FaceDetectionOptions): Promise<DetectedFaces>;
 }

server/src/interfaces/person.interface.ts

@@ -1,9 +1,10 @@
+import { Insertable, Updateable } from 'kysely';
+import { AssetFaces, FaceSearch, Person } from 'src/db';
 import { AssetFaceEntity } from 'src/entities/asset-face.entity';
-import { FaceSearchEntity } from 'src/entities/face-search.entity';
 import { PersonEntity } from 'src/entities/person.entity';
 import { SourceType } from 'src/enum';
 import { Paginated, PaginationOptions } from 'src/utils/pagination';
-import { FindManyOptions, FindOptionsRelations, FindOptionsSelect } from 'typeorm';
+import { FindOptionsRelations } from 'typeorm';
 
 export const IPersonRepository = 'IPersonRepository';
 
@@ -48,29 +49,31 @@ export interface DeleteFacesOptions {
 
 export type UnassignFacesOptions = DeleteFacesOptions;
 
+export type SelectFaceOptions = Partial<{ [K in keyof AssetFaceEntity]: boolean }>;
+
 export interface IPersonRepository {
-  getAll(pagination: PaginationOptions, options?: FindManyOptions<PersonEntity>): Paginated<PersonEntity>;
+  getAll(options?: Partial<PersonEntity>): AsyncIterableIterator<PersonEntity>;
   getAllForUser(pagination: PaginationOptions, userId: string, options: PersonSearchOptions): Paginated<PersonEntity>;
   getAllWithoutFaces(): Promise<PersonEntity[]>;
   getById(personId: string): Promise<PersonEntity | null>;
   getByName(userId: string, personName: string, options: PersonNameSearchOptions): Promise<PersonEntity[]>;
   getDistinctNames(userId: string, options: PersonNameSearchOptions): Promise<PersonNameResponse[]>;
 
-  create(person: Partial<PersonEntity>): Promise<PersonEntity>;
-  createAll(people: Partial<PersonEntity>[]): Promise<string[]>;
+  create(person: Insertable<Person>): Promise<PersonEntity>;
+  createAll(people: Insertable<Person>[]): Promise<string[]>;
   delete(entities: PersonEntity[]): Promise<void>;
   deleteFaces(options: DeleteFacesOptions): Promise<void>;
   refreshFaces(
-    facesToAdd: Partial<AssetFaceEntity>[],
+    facesToAdd: Insertable<AssetFaces>[],
     faceIdsToRemove: string[],
-    embeddingsToAdd?: FaceSearchEntity[],
+    embeddingsToAdd?: Insertable<FaceSearch>[],
   ): Promise<void>;
-  getAllFaces(pagination: PaginationOptions, options?: FindManyOptions<AssetFaceEntity>): Paginated<AssetFaceEntity>;
+  getAllFaces(options?: Partial<AssetFaceEntity>): AsyncIterableIterator<AssetFaceEntity>;
   getFaceById(id: string): Promise<AssetFaceEntity>;
   getFaceByIdWithAssets(
     id: string,
     relations?: FindOptionsRelations<AssetFaceEntity>,
-    select?: FindOptionsSelect<AssetFaceEntity>,
+    select?: SelectFaceOptions,
   ): Promise<AssetFaceEntity | null>;
   getFaces(assetId: string): Promise<AssetFaceEntity[]>;
   getFacesByIds(ids: AssetFaceId[]): Promise<AssetFaceEntity[]>;
@@ -80,7 +83,7 @@ export interface IPersonRepository {
   getNumberOfPeople(userId: string): Promise<PeopleStatistics>;
   reassignFaces(data: UpdateFacesData): Promise<number>;
   unassignFaces(options: UnassignFacesOptions): Promise<void>;
-  update(person: Partial<PersonEntity>): Promise<PersonEntity>;
-  updateAll(people: Partial<PersonEntity>[]): Promise<void>;
+  update(person: Updateable<Person> & { id: string }): Promise<PersonEntity>;
+  updateAll(people: Insertable<Person>[]): Promise<void>;
   getLatestFaceDate(): Promise<string | undefined>;
 }

server/src/interfaces/search.interface.ts

@@ -104,7 +104,7 @@ export interface SearchExifOptions {
 }
 
 export interface SearchEmbeddingOptions {
-  embedding: number[];
+  embedding: string;
   userIds: string[];
 }
 
@@ -152,7 +152,7 @@ export interface FaceEmbeddingSearch extends SearchEmbeddingOptions {
 
 export interface AssetDuplicateSearch {
   assetId: string;
-  embedding: number[];
+  embedding: string;
   maxDistance: number;
   type: AssetType;
   userIds: string[];
@@ -192,7 +192,7 @@ export interface ISearchRepository {
   searchDuplicates(options: AssetDuplicateSearch): Promise<AssetDuplicateResult[]>;
   searchFaces(search: FaceEmbeddingSearch): Promise<FaceSearchResult[]>;
   searchRandom(size: number, options: AssetSearchOptions): Promise<AssetEntity[]>;
-  upsert(assetId: string, embedding: number[]): Promise<void>;
+  upsert(assetId: string, embedding: string): Promise<void>;
   searchPlaces(placeName: string): Promise<GeodataPlacesEntity[]>;
   getAssetsByCity(userIds: string[]): Promise<AssetEntity[]>;
   deleteAllSearchEmbeddings(): Promise<void>;