From 6b1b5054f852641b7b7626b459fbc393d760dee9 Mon Sep 17 00:00:00 2001 From: Mert <101130780+mertalev@users.noreply.github.com> Date: Sun, 16 Jun 2024 15:25:27 -0400 Subject: [PATCH] feat(server): separate face search relation (#10371) * wip * various fixes * new migration * fix test * add face search entity, update sql * update e2e * set storage to external --- e2e/src/utils.ts | 9 +--- server/src/entities/asset-face.entity.ts | 8 +-- server/src/entities/face-search.entity.ts | 21 ++++++++ server/src/entities/index.ts | 2 + .../1718486162779-AddFaceSearchRelation.ts | 54 +++++++++++++++++++ server/src/queries/search.repository.sql | 5 +- .../src/repositories/database.repository.ts | 2 +- server/src/repositories/person.repository.ts | 7 +-- server/src/repositories/search.repository.ts | 5 +- server/src/services/person.service.spec.ts | 5 +- server/src/services/person.service.ts | 41 ++++++++------ server/test/fixtures/face.stub.ts | 18 +++---- 12 files changed, 130 insertions(+), 47 deletions(-) create mode 100644 server/src/entities/face-search.entity.ts create mode 100644 server/src/migrations/1718486162779-AddFaceSearchRelation.ts diff --git a/e2e/src/utils.ts b/e2e/src/utils.ts index 72b3480299..f3afdddf8b 100644 --- a/e2e/src/utils.ts +++ b/e2e/src/utils.ts @@ -398,14 +398,7 @@ export const utils = { return; } - const vector = Array.from({ length: 512 }, Math.random); - const embedding = `[${vector.join(',')}]`; - - await client.query('INSERT INTO asset_faces ("assetId", "personId", "embedding") VALUES ($1, $2, $3)', [ - assetId, - personId, - embedding, - ]); + await client.query('INSERT INTO asset_faces ("assetId", "personId") VALUES ($1, $2)', [assetId, personId]); }, setPersonThumbnail: async (personId: string) => { diff --git a/server/src/entities/asset-face.entity.ts b/server/src/entities/asset-face.entity.ts index 38fcd46063..c21aacfcd1 100644 --- a/server/src/entities/asset-face.entity.ts +++ b/server/src/entities/asset-face.entity.ts @@ -1,6 +1,7 @@ import { AssetEntity } from 'src/entities/asset.entity'; +import { FaceSearchEntity } from 'src/entities/face-search.entity'; import { PersonEntity } from 'src/entities/person.entity'; -import { Column, Entity, Index, ManyToOne, PrimaryGeneratedColumn } from 'typeorm'; +import { Column, Entity, Index, ManyToOne, OneToOne, PrimaryGeneratedColumn } from 'typeorm'; @Entity('asset_faces', { synchronize: false }) @Index('IDX_asset_faces_assetId_personId', ['assetId', 'personId']) @@ -15,9 +16,8 @@ export class AssetFaceEntity { @Column({ nullable: true, type: 'uuid' }) personId!: string | null; - @Index('face_index', { synchronize: false }) - @Column({ type: 'float4', array: true, select: false, transformer: { from: (v) => JSON.parse(v), to: (v) => v } }) - embedding!: number[]; + @OneToOne(() => FaceSearchEntity, (faceSearchEntity) => faceSearchEntity.face, { cascade: ['insert'] }) + faceSearch?: FaceSearchEntity; @Column({ default: 0, type: 'int' }) imageWidth!: number; diff --git a/server/src/entities/face-search.entity.ts b/server/src/entities/face-search.entity.ts new file mode 100644 index 0000000000..3fd3c65f28 --- /dev/null +++ b/server/src/entities/face-search.entity.ts @@ -0,0 +1,21 @@ +import { AssetFaceEntity } from 'src/entities/asset-face.entity'; +import { asVector } from 'src/utils/database'; +import { Column, Entity, Index, JoinColumn, OneToOne, PrimaryColumn } from 'typeorm'; + +@Entity('face_search', { synchronize: false }) +export class FaceSearchEntity { + @OneToOne(() => AssetFaceEntity, { onDelete: 'CASCADE', nullable: true }) + @JoinColumn({ name: 'faceId', referencedColumnName: 'id' }) + face?: AssetFaceEntity; + + @PrimaryColumn() + faceId!: string; + + @Index('face_index', { synchronize: false }) + @Column({ + type: 'float4', + array: true, + transformer: { from: (v) => JSON.parse(v), to: (v) => asVector(v) }, + }) + embedding!: number[]; +} diff --git a/server/src/entities/index.ts b/server/src/entities/index.ts index 313f2dc269..cd3d74724b 100644 --- a/server/src/entities/index.ts +++ b/server/src/entities/index.ts @@ -8,6 +8,7 @@ import { AssetStackEntity } from 'src/entities/asset-stack.entity'; import { AssetEntity } from 'src/entities/asset.entity'; import { AuditEntity } from 'src/entities/audit.entity'; import { ExifEntity } from 'src/entities/exif.entity'; +import { FaceSearchEntity } from 'src/entities/face-search.entity'; import { GeodataPlacesEntity } from 'src/entities/geodata-places.entity'; import { LibraryEntity } from 'src/entities/library.entity'; import { MemoryEntity } from 'src/entities/memory.entity'; @@ -34,6 +35,7 @@ export const entities = [ AssetJobStatusEntity, AuditEntity, ExifEntity, + FaceSearchEntity, GeodataPlacesEntity, MemoryEntity, MoveEntity, diff --git a/server/src/migrations/1718486162779-AddFaceSearchRelation.ts b/server/src/migrations/1718486162779-AddFaceSearchRelation.ts new file mode 100644 index 0000000000..5bf3fcd97b --- /dev/null +++ b/server/src/migrations/1718486162779-AddFaceSearchRelation.ts @@ -0,0 +1,54 @@ +import { getVectorExtension } from 'src/database.config'; +import { DatabaseExtension } from 'src/interfaces/database.interface'; +import { MigrationInterface, QueryRunner } from 'typeorm'; + +export class AddFaceSearchRelation1718486162779 implements MigrationInterface { + public async up(queryRunner: QueryRunner): Promise { + if (getVectorExtension() === DatabaseExtension.VECTORS) { + await queryRunner.query(`SET search_path TO "$user", public, vectors`); + await queryRunner.query(`SET vectors.pgvector_compatibility=on`); + } + + await queryRunner.query(` + CREATE TABLE face_search ( + "faceId" uuid PRIMARY KEY REFERENCES asset_faces(id) ON DELETE CASCADE, + embedding vector(512) NOT NULL )`); + + await queryRunner.query(`ALTER TABLE face_search ALTER COLUMN embedding SET STORAGE EXTERNAL`); + await queryRunner.query(`ALTER TABLE smart_search ALTER COLUMN embedding SET STORAGE EXTERNAL`); + + await queryRunner.query(` + INSERT INTO face_search("faceId", embedding) + SELECT id, embedding + FROM asset_faces faces`); + + await queryRunner.query(`ALTER TABLE asset_faces DROP COLUMN "embedding"`); + + await queryRunner.query(` + CREATE INDEX face_index ON face_search + USING hnsw (embedding vector_cosine_ops) + WITH (ef_construction = 300, m = 16)`); + } + + public async down(queryRunner: QueryRunner): Promise { + if (getVectorExtension() === DatabaseExtension.VECTORS) { + await queryRunner.query(`SET search_path TO "$user", public, vectors`); + await queryRunner.query(`SET vectors.pgvector_compatibility=on`); + } + + await queryRunner.query(`ALTER TABLE asset_faces ADD COLUMN "embedding" vector(512)`); + await queryRunner.query(`ALTER TABLE face_search ALTER COLUMN embedding SET STORAGE DEFAULT`); + await queryRunner.query(`ALTER TABLE smart_search ALTER COLUMN embedding SET STORAGE DEFAULT`); + await queryRunner.query(` + UPDATE asset_faces + SET embedding = fs.embedding + FROM face_search fs + WHERE id = fs."faceId"`); + await queryRunner.query(`DROP TABLE face_search`); + + await queryRunner.query(` + CREATE INDEX face_index ON asset_faces + USING hnsw (embedding vector_cosine_ops) + WITH (ef_construction = 300, m = 16)`); + } +} diff --git a/server/src/queries/search.repository.sql b/server/src/queries/search.repository.sql index 9efeae6248..987828a860 100644 --- a/server/src/queries/search.repository.sql +++ b/server/src/queries/search.repository.sql @@ -241,15 +241,16 @@ WITH "faces"."boundingBoxY1" AS "boundingBoxY1", "faces"."boundingBoxX2" AS "boundingBoxX2", "faces"."boundingBoxY2" AS "boundingBoxY2", - "faces"."embedding" <= > $1 AS "distance" + "search"."embedding" <= > $1 AS "distance" FROM "asset_faces" "faces" INNER JOIN "assets" "asset" ON "asset"."id" = "faces"."assetId" AND ("asset"."deletedAt" IS NULL) + INNER JOIN "face_search" "search" ON "search"."faceId" = "faces"."id" WHERE "asset"."ownerId" IN ($2) ORDER BY - "faces"."embedding" <= > $1 ASC + "search"."embedding" <= > $1 ASC LIMIT 100 ) diff --git a/server/src/repositories/database.repository.ts b/server/src/repositories/database.repository.ts index dc442e7017..fc9e76b0aa 100644 --- a/server/src/repositories/database.repository.ts +++ b/server/src/repositories/database.repository.ts @@ -98,7 +98,7 @@ export class DatabaseRepository implements IDatabaseRepository { } catch (error) { if (getVectorExtension() === DatabaseExtension.VECTORS) { this.logger.warn(`Could not reindex index ${index}. Attempting to auto-fix.`); - const table = index === VectorIndex.CLIP ? 'smart_search' : 'asset_faces'; + const table = index === VectorIndex.CLIP ? 'smart_search' : 'face_search'; const dimSize = await this.getDimSize(table); await this.dataSource.manager.transaction(async (manager) => { await this.setSearchPath(manager); diff --git a/server/src/repositories/person.repository.ts b/server/src/repositories/person.repository.ts index 225a2edeca..36d742f8dc 100644 --- a/server/src/repositories/person.repository.ts +++ b/server/src/repositories/person.repository.ts @@ -14,7 +14,6 @@ import { PersonStatistics, UpdateFacesData, } from 'src/interfaces/person.interface'; -import { asVector } from 'src/utils/database'; import { Instrumentation } from 'src/utils/instrumentation'; import { Paginated, PaginationOptions, paginate } from 'src/utils/pagination'; import { FindManyOptions, FindOptionsRelations, FindOptionsSelect, In, Repository } from 'typeorm'; @@ -249,10 +248,8 @@ export class PersonRepository implements IPersonRepository { } async createFaces(entities: AssetFaceEntity[]): Promise { - const res = await this.assetFaceRepository.insert( - entities.map((entity) => ({ ...entity, embedding: () => asVector(entity.embedding, true) })), - ); - return res.identifiers.map((row) => row.id); + const res = await this.assetFaceRepository.save(entities); + return res.map((row) => row.id); } async update(entity: Partial): Promise { diff --git a/server/src/repositories/search.repository.ts b/server/src/repositories/search.repository.ts index f0c5dcb364..439ccd099c 100644 --- a/server/src/repositories/search.repository.ts +++ b/server/src/repositories/search.repository.ts @@ -218,10 +218,11 @@ export class SearchRepository implements ISearchRepository { await this.assetRepository.manager.transaction(async (manager) => { const cte = manager .createQueryBuilder(AssetFaceEntity, 'faces') - .select('faces.embedding <=> :embedding', 'distance') + .select('search.embedding <=> :embedding', 'distance') .innerJoin('faces.asset', 'asset') + .innerJoin('faces.faceSearch', 'search') .where('asset.ownerId IN (:...userIds )') - .orderBy('faces.embedding <=> :embedding') + .orderBy('search.embedding <=> :embedding') .setParameters({ userIds, embedding: asVector(embedding) }); cte.limit(numResults); diff --git a/server/src/services/person.service.spec.ts b/server/src/services/person.service.spec.ts index bf6fc8207e..eb0e3ad1e9 100644 --- a/server/src/services/person.service.spec.ts +++ b/server/src/services/person.service.spec.ts @@ -668,15 +668,18 @@ describe(PersonService.name, () => { machineLearningMock.detectFaces.mockResolvedValue(detectFaceMock); searchMock.searchFaces.mockResolvedValue([{ face: faceStub.face1, distance: 0.7 }]); assetMock.getByIds.mockResolvedValue([assetStub.image]); + const faceId = 'face-id'; + cryptoMock.randomUUID.mockReturnValue(faceId); const face = { + id: faceId, assetId: 'asset-id', - embedding: [1, 2, 3, 4], boundingBoxX1: 100, boundingBoxY1: 100, boundingBoxX2: 200, boundingBoxY2: 200, imageHeight: 500, imageWidth: 400, + faceSearch: { faceId, embedding: [1, 2, 3, 4] }, }; await sut.handleDetectFaces({ id: assetStub.image.id }); diff --git a/server/src/services/person.service.ts b/server/src/services/person.service.ts index 57940f3113..05034dc6f9 100644 --- a/server/src/services/person.service.ts +++ b/server/src/services/person.service.ts @@ -22,6 +22,7 @@ import { mapFaces, mapPerson, } from 'src/dtos/person.dto'; +import { AssetFaceEntity } from 'src/entities/asset-face.entity'; import { AssetEntity, AssetType } from 'src/entities/asset.entity'; import { PersonPathType } from 'src/entities/move.entity'; import { PersonEntity } from 'src/entities/person.entity'; @@ -70,7 +71,7 @@ export class PersonService { @Inject(IStorageRepository) private storageRepository: IStorageRepository, @Inject(IJobRepository) private jobRepository: IJobRepository, @Inject(ISearchRepository) private smartInfoRepository: ISearchRepository, - @Inject(ICryptoRepository) cryptoRepository: ICryptoRepository, + @Inject(ICryptoRepository) private cryptoRepository: ICryptoRepository, @Inject(ILoggerRepository) private logger: ILoggerRepository, ) { this.access = AccessCore.create(accessRepository); @@ -347,16 +348,21 @@ export class PersonService { if (faces.length > 0) { await this.jobRepository.queue({ name: JobName.QUEUE_FACIAL_RECOGNITION, data: { force: false } }); - const mappedFaces = faces.map((face) => ({ - assetId: asset.id, - embedding: face.embedding, - imageHeight, - imageWidth, - boundingBoxX1: face.boundingBox.x1, - boundingBoxY1: face.boundingBox.y1, - boundingBoxX2: face.boundingBox.x2, - boundingBoxY2: face.boundingBox.y2, - })); + const mappedFaces: Partial[] = []; + for (const face of faces) { + const faceId = this.cryptoRepository.randomUUID(); + mappedFaces.push({ + id: faceId, + assetId: asset.id, + imageHeight, + imageWidth, + boundingBoxX1: face.boundingBox.x1, + boundingBoxY1: face.boundingBox.y1, + boundingBoxX2: face.boundingBox.x2, + boundingBoxY2: face.boundingBox.y2, + faceSearch: { faceId, embedding: face.embedding }, + }); + } const faceIds = await this.repository.createFaces(mappedFaces); await this.jobRepository.queueAll(faceIds.map((id) => ({ name: JobName.FACIAL_RECOGNITION, data: { id } }))); @@ -409,14 +415,19 @@ export class PersonService { const face = await this.repository.getFaceByIdWithAssets( id, - { person: true, asset: true }, - { id: true, personId: true, embedding: true }, + { person: true, asset: true, faceSearch: true }, + { id: true, personId: true, faceSearch: { embedding: true } }, ); if (!face || !face.asset) { this.logger.warn(`Face ${id} not found`); return JobStatus.FAILED; } + if (!face.faceSearch?.embedding) { + this.logger.warn(`Face ${id} does not have an embedding`); + return JobStatus.FAILED; + } + if (face.personId) { this.logger.debug(`Face ${id} already has a person assigned`); return JobStatus.SKIPPED; @@ -424,7 +435,7 @@ export class PersonService { const matches = await this.smartInfoRepository.searchFaces({ userIds: [face.asset.ownerId], - embedding: face.embedding, + embedding: face.faceSearch.embedding, maxDistance: machineLearning.facialRecognition.maxDistance, numResults: machineLearning.facialRecognition.minFaces, }); @@ -448,7 +459,7 @@ export class PersonService { if (!personId) { const matchWithPerson = await this.smartInfoRepository.searchFaces({ userIds: [face.asset.ownerId], - embedding: face.embedding, + embedding: face.faceSearch.embedding, maxDistance: machineLearning.facialRecognition.maxDistance, numResults: 1, hasPerson: true, diff --git a/server/test/fixtures/face.stub.ts b/server/test/fixtures/face.stub.ts index 5ecb5701ce..82935dd345 100644 --- a/server/test/fixtures/face.stub.ts +++ b/server/test/fixtures/face.stub.ts @@ -11,13 +11,13 @@ export const faceStub = { asset: assetStub.image, personId: personStub.withName.id, person: personStub.withName, - embedding: [1, 2, 3, 4], boundingBoxX1: 0, boundingBoxY1: 0, boundingBoxX2: 1, boundingBoxY2: 1, imageHeight: 1024, imageWidth: 1024, + faceSearch: { faceId: 'assetFaceId1', embedding: [1, 2, 3, 4] }, }), primaryFace1: Object.freeze>({ id: 'assetFaceId2', @@ -25,13 +25,13 @@ export const faceStub = { asset: assetStub.image, personId: personStub.primaryPerson.id, person: personStub.primaryPerson, - embedding: [1, 2, 3, 4], boundingBoxX1: 0, boundingBoxY1: 0, boundingBoxX2: 1, boundingBoxY2: 1, imageHeight: 1024, imageWidth: 1024, + faceSearch: { faceId: 'assetFaceId2', embedding: [1, 2, 3, 4] }, }), mergeFace1: Object.freeze>({ id: 'assetFaceId3', @@ -39,13 +39,13 @@ export const faceStub = { asset: assetStub.image, personId: personStub.mergePerson.id, person: personStub.mergePerson, - embedding: [1, 2, 3, 4], boundingBoxX1: 0, boundingBoxY1: 0, boundingBoxX2: 1, boundingBoxY2: 1, imageHeight: 1024, imageWidth: 1024, + faceSearch: { faceId: 'assetFaceId3', embedding: [1, 2, 3, 4] }, }), mergeFace2: Object.freeze>({ id: 'assetFaceId4', @@ -53,13 +53,13 @@ export const faceStub = { asset: assetStub.image1, personId: personStub.mergePerson.id, person: personStub.mergePerson, - embedding: [1, 2, 3, 4], boundingBoxX1: 0, boundingBoxY1: 0, boundingBoxX2: 1, boundingBoxY2: 1, imageHeight: 1024, imageWidth: 1024, + faceSearch: { faceId: 'assetFaceId4', embedding: [1, 2, 3, 4] }, }), start: Object.freeze>({ id: 'assetFaceId5', @@ -67,13 +67,13 @@ export const faceStub = { asset: assetStub.image, personId: personStub.newThumbnail.id, person: personStub.newThumbnail, - embedding: [1, 2, 3, 4], boundingBoxX1: 5, boundingBoxY1: 5, boundingBoxX2: 505, boundingBoxY2: 505, imageHeight: 2880, imageWidth: 2160, + faceSearch: { faceId: 'assetFaceId5', embedding: [1, 2, 3, 4] }, }), middle: Object.freeze>({ id: 'assetFaceId6', @@ -81,13 +81,13 @@ export const faceStub = { asset: assetStub.image, personId: personStub.newThumbnail.id, person: personStub.newThumbnail, - embedding: [1, 2, 3, 4], boundingBoxX1: 100, boundingBoxY1: 100, boundingBoxX2: 200, boundingBoxY2: 200, imageHeight: 500, imageWidth: 400, + faceSearch: { faceId: 'assetFaceId6', embedding: [1, 2, 3, 4] }, }), end: Object.freeze>({ id: 'assetFaceId7', @@ -95,13 +95,13 @@ export const faceStub = { asset: assetStub.image, personId: personStub.newThumbnail.id, person: personStub.newThumbnail, - embedding: [1, 2, 3, 4], boundingBoxX1: 300, boundingBoxY1: 300, boundingBoxX2: 495, boundingBoxY2: 495, imageHeight: 500, imageWidth: 500, + faceSearch: { faceId: 'assetFaceId7', embedding: [1, 2, 3, 4] }, }), noPerson1: Object.freeze({ id: 'assetFaceId8', @@ -109,13 +109,13 @@ export const faceStub = { asset: assetStub.image, personId: null, person: null, - embedding: [1, 2, 3, 4], boundingBoxX1: 0, boundingBoxY1: 0, boundingBoxX2: 1, boundingBoxY2: 1, imageHeight: 1024, imageWidth: 1024, + faceSearch: { faceId: 'assetFaceId8', embedding: [1, 2, 3, 4] }, }), noPerson2: Object.freeze({ id: 'assetFaceId9', @@ -123,12 +123,12 @@ export const faceStub = { asset: assetStub.image, personId: null, person: null, - embedding: [1, 2, 3, 4], boundingBoxX1: 0, boundingBoxY1: 0, boundingBoxX2: 1, boundingBoxY2: 1, imageHeight: 1024, imageWidth: 1024, + faceSearch: { faceId: 'assetFaceId9', embedding: [1, 2, 3, 4] }, }), };