1
0
mirror of https://github.com/immich-app/immich.git synced 2024-11-24 08:52:28 +02:00

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
This commit is contained in:
Mert 2024-06-16 15:25:27 -04:00 committed by GitHub
parent 0fe152b1ef
commit 6b1b5054f8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
12 changed files with 130 additions and 47 deletions

View File

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

View File

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

View File

@ -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[];
}

View File

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

View File

@ -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<void> {
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<void> {
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)`);
}
}

View File

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

View File

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

View File

@ -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<string[]> {
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<PersonEntity>): Promise<PersonEntity> {

View File

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

View File

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

View File

@ -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) => ({
const mappedFaces: Partial<AssetFaceEntity>[] = [];
for (const face of faces) {
const faceId = this.cryptoRepository.randomUUID();
mappedFaces.push({
id: faceId,
assetId: asset.id,
embedding: face.embedding,
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,

View File

@ -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<NonNullableProperty<AssetFaceEntity>>({
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<NonNullableProperty<AssetFaceEntity>>({
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<NonNullableProperty<AssetFaceEntity>>({
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<NonNullableProperty<AssetFaceEntity>>({
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<NonNullableProperty<AssetFaceEntity>>({
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<NonNullableProperty<AssetFaceEntity>>({
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<AssetFaceEntity>({
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<AssetFaceEntity>({
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] },
}),
};