feat: facial recognition (#2180)

2023-05-17 13:07:17 -04:00
parent 115a47d4c6
commit 93863b0629
107 changed files with 3943 additions and 133 deletions
@@ -1,8 +1,9 @@
-import { AlbumEntity, AssetEntity, AssetType } from '@app/infra/entities';
+import { AlbumEntity, AssetEntity, AssetFaceEntity, AssetType } from '@app/infra/entities';

 export enum SearchCollection {
  ASSETS = 'assets',
  ALBUMS = 'albums',
+  FACES = 'faces',
 }

 export enum SearchStrategy {
@@ -10,6 +11,10 @@ export enum SearchStrategy {
  TEXT = 'TEXT',
 }

+export interface SearchFaceFilter {
+  ownerId: string;
+}
+
 export interface SearchFilter {
  id?: string;
  userId: string;
@@ -37,6 +42,8 @@ export interface SearchResult<T> {
  page: number;
  /** items for page */
  items: T[];
+  /** score */
+  distances: number[];
  facets: SearchFacet[];
 }

@@ -56,6 +63,13 @@ export interface SearchExploreItem<T> {
  }>;
 }

+export type OwnedFaceEntity = Pick<AssetFaceEntity, 'assetId' | 'personId' | 'embedding'> & {
+  /** computed as assetId|personId */
+  id: string;
+  /** copied from asset.id */
+  ownerId: string;
+};
+
 export type SearchCollectionIndexStatus = Record<SearchCollection, boolean>;

 export const ISearchRepository = 'ISearchRepository';
@@ -66,13 +80,17 @@ export interface ISearchRepository {

  importAlbums(items: AlbumEntity[], done: boolean): Promise<void>;
  importAssets(items: AssetEntity[], done: boolean): Promise<void>;
+  importFaces(items: OwnedFaceEntity[], done: boolean): Promise<void>;

  deleteAlbums(ids: string[]): Promise<void>;
  deleteAssets(ids: string[]): Promise<void>;
+  deleteFaces(ids: string[]): Promise<void>;
+  deleteAllFaces(): Promise<number>;

  searchAlbums(query: string, filters: SearchFilter): Promise<SearchResult<AlbumEntity>>;
  searchAssets(query: string, filters: SearchFilter): Promise<SearchResult<AssetEntity>>;
  vectorSearch(query: number[], filters: SearchFilter): Promise<SearchResult<AssetEntity>>;
+  searchFaces(query: number[], filters: SearchFaceFilter): Promise<SearchResult<AssetFaceEntity>>;

  explore(userId: string): Promise<SearchExploreItem<AssetEntity>[]>;
 }
@@ -6,8 +6,10 @@ import {
  assetEntityStub,
  asyncTick,
  authStub,
+  faceStub,
  newAlbumRepositoryMock,
  newAssetRepositoryMock,
+  newFaceRepositoryMock,
  newJobRepositoryMock,
  newMachineLearningRepositoryMock,
  newSearchRepositoryMock,
@@ -15,6 +17,7 @@ import {
 } from '../../test';
 import { IAlbumRepository } from '../album/album.repository';
 import { IAssetRepository } from '../asset/asset.repository';
+import { IFaceRepository } from '../facial-recognition';
 import { JobName } from '../job';
 import { IJobRepository } from '../job/job.repository';
 import { IMachineLearningRepository } from '../smart-info';
@@ -28,20 +31,29 @@ describe(SearchService.name, () => {
  let sut: SearchService;
  let albumMock: jest.Mocked<IAlbumRepository>;
  let assetMock: jest.Mocked<IAssetRepository>;
+  let faceMock: jest.Mocked<IFaceRepository>;
  let jobMock: jest.Mocked<IJobRepository>;
  let machineMock: jest.Mocked<IMachineLearningRepository>;
  let searchMock: jest.Mocked<ISearchRepository>;
  let configMock: jest.Mocked<ConfigService>;

+  const makeSut = (value?: string) => {
+    if (value) {
+      configMock.get.mockReturnValue(value);
+    }
+    return new SearchService(albumMock, assetMock, faceMock, jobMock, machineMock, searchMock, configMock);
+  };
+
  beforeEach(() => {
    albumMock = newAlbumRepositoryMock();
    assetMock = newAssetRepositoryMock();
+    faceMock = newFaceRepositoryMock();
    jobMock = newJobRepositoryMock();
    machineMock = newMachineLearningRepositoryMock();
    searchMock = newSearchRepositoryMock();
    configMock = { get: jest.fn() } as unknown as jest.Mocked<ConfigService>;

-    sut = new SearchService(albumMock, assetMock, jobMock, machineMock, searchMock, configMock);
+    sut = makeSut();
  });

  afterEach(() => {
@@ -80,8 +92,7 @@ describe(SearchService.name, () => {
    });

    it('should be disabled via an env variable', () => {
-      configMock.get.mockReturnValue('false');
-      const sut = new SearchService(albumMock, assetMock, jobMock, machineMock, searchMock, configMock);
+      const sut = makeSut('false');

      expect(sut.isEnabled()).toBe(false);
    });
@@ -93,8 +104,7 @@ describe(SearchService.name, () => {
    });

    it('should return the config when search is disabled', () => {
-      configMock.get.mockReturnValue('false');
-      const sut = new SearchService(albumMock, assetMock, jobMock, machineMock, searchMock, configMock);
+      const sut = makeSut('false');

      expect(sut.getConfig()).toEqual({ enabled: false });
    });
@@ -102,8 +112,7 @@ describe(SearchService.name, () => {

  describe(`bootstrap`, () => {
    it('should skip when search is disabled', async () => {
-      configMock.get.mockReturnValue('false');
-      const sut = new SearchService(albumMock, assetMock, jobMock, machineMock, searchMock, configMock);
+      const sut = makeSut('false');

      await sut.bootstrap();

@@ -115,7 +124,7 @@ describe(SearchService.name, () => {
    });

    it('should skip schema migration if not needed', async () => {
-      searchMock.checkMigrationStatus.mockResolvedValue({ assets: false, albums: false });
+      searchMock.checkMigrationStatus.mockResolvedValue({ assets: false, albums: false, faces: false });
      await sut.bootstrap();

      expect(searchMock.setup).toHaveBeenCalled();
@@ -123,21 +132,21 @@ describe(SearchService.name, () => {
    });

    it('should do schema migration if needed', async () => {
-      searchMock.checkMigrationStatus.mockResolvedValue({ assets: true, albums: true });
+      searchMock.checkMigrationStatus.mockResolvedValue({ assets: true, albums: true, faces: true });
      await sut.bootstrap();

      expect(searchMock.setup).toHaveBeenCalled();
      expect(jobMock.queue.mock.calls).toEqual([
        [{ name: JobName.SEARCH_INDEX_ASSETS }],
        [{ name: JobName.SEARCH_INDEX_ALBUMS }],
+        [{ name: JobName.SEARCH_INDEX_FACES }],
      ]);
    });
  });

  describe('search', () => {
    it('should throw an error is search is disabled', async () => {
-      configMock.get.mockReturnValue('false');
-      const sut = new SearchService(albumMock, assetMock, jobMock, machineMock, searchMock, configMock);
+      const sut = makeSut('false');

      await expect(sut.search(authStub.admin, {})).rejects.toBeInstanceOf(BadRequestException);

@@ -157,6 +166,7 @@ describe(SearchService.name, () => {
          page: 1,
          items: [],
          facets: [],
+          distances: [],
        },
        assets: {
          total: 0,
@@ -164,6 +174,7 @@ describe(SearchService.name, () => {
          page: 1,
          items: [],
          facets: [],
+          distances: [],
        },
      });

@@ -202,8 +213,7 @@ describe(SearchService.name, () => {
    });

    it('should skip if search is disabled', async () => {
-      configMock.get.mockReturnValue('false');
-      const sut = new SearchService(albumMock, assetMock, jobMock, machineMock, searchMock, configMock);
+      const sut = makeSut('false');

      await sut.handleIndexAssets();

@@ -214,8 +224,7 @@ describe(SearchService.name, () => {

  describe('handleIndexAsset', () => {
    it('should skip if search is disabled', () => {
-      configMock.get.mockReturnValue('false');
-      const sut = new SearchService(albumMock, assetMock, jobMock, machineMock, searchMock, configMock);
+      const sut = makeSut('false');
      sut.handleIndexAsset({ ids: [assetEntityStub.image.id] });
    });

@@ -226,8 +235,7 @@ describe(SearchService.name, () => {

  describe('handleIndexAlbums', () => {
    it('should skip if search is disabled', () => {
-      configMock.get.mockReturnValue('false');
-      const sut = new SearchService(albumMock, assetMock, jobMock, machineMock, searchMock, configMock);
+      const sut = makeSut('false');
      sut.handleIndexAlbums();
    });

@@ -251,8 +259,7 @@ describe(SearchService.name, () => {

  describe('handleIndexAlbum', () => {
    it('should skip if search is disabled', () => {
-      configMock.get.mockReturnValue('false');
-      const sut = new SearchService(albumMock, assetMock, jobMock, machineMock, searchMock, configMock);
+      const sut = makeSut('false');
      sut.handleIndexAlbum({ ids: [albumStub.empty.id] });
    });

@@ -263,8 +270,7 @@ describe(SearchService.name, () => {

  describe('handleRemoveAlbum', () => {
    it('should skip if search is disabled', () => {
-      configMock.get.mockReturnValue('false');
-      const sut = new SearchService(albumMock, assetMock, jobMock, machineMock, searchMock, configMock);
+      const sut = makeSut('false');
      sut.handleRemoveAlbum({ ids: ['album1'] });
    });

@@ -275,8 +281,7 @@ describe(SearchService.name, () => {

  describe('handleRemoveAsset', () => {
    it('should skip if search is disabled', () => {
-      configMock.get.mockReturnValue('false');
-      const sut = new SearchService(albumMock, assetMock, jobMock, machineMock, searchMock, configMock);
+      const sut = makeSut('false');
      sut.handleRemoveAsset({ ids: ['asset1'] });
    });

@@ -285,6 +290,84 @@ describe(SearchService.name, () => {
    });
  });

+  describe('handleIndexFaces', () => {
+    it('should call done, even when there are no faces', async () => {
+      faceMock.getAll.mockResolvedValue([]);
+
+      await sut.handleIndexFaces();
+
+      expect(searchMock.importFaces).toHaveBeenCalledWith([], true);
+    });
+
+    it('should index all the faces', async () => {
+      faceMock.getAll.mockResolvedValue([faceStub.face1]);
+
+      await sut.handleIndexFaces();
+
+      expect(searchMock.importFaces.mock.calls).toEqual([
+        [
+          [
+            {
+              id: 'asset-id|person-1',
+              ownerId: 'user-id',
+              assetId: 'asset-id',
+              personId: 'person-1',
+              embedding: [1, 2, 3, 4],
+            },
+          ],
+          false,
+        ],
+        [[], true],
+      ]);
+    });
+
+    it('should log an error', async () => {
+      faceMock.getAll.mockResolvedValue([faceStub.face1]);
+      searchMock.importFaces.mockRejectedValue(new Error('import failed'));
+
+      await sut.handleIndexFaces();
+
+      expect(searchMock.importFaces).toHaveBeenCalled();
+    });
+
+    it('should skip if search is disabled', async () => {
+      const sut = makeSut('false');
+
+      await sut.handleIndexFaces();
+
+      expect(searchMock.importFaces).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('handleIndexAsset', () => {
+    it('should skip if search is disabled', () => {
+      const sut = makeSut('false');
+      sut.handleIndexFace({ assetId: 'asset-1', personId: 'person-1' });
+
+      expect(searchMock.importFaces).not.toHaveBeenCalled();
+      expect(faceMock.getByIds).not.toHaveBeenCalled();
+    });
+
+    it('should index the face', () => {
+      faceMock.getByIds.mockResolvedValue([faceStub.face1]);
+
+      sut.handleIndexFace({ assetId: 'asset-1', personId: 'person-1' });
+
+      expect(faceMock.getByIds).toHaveBeenCalledWith([{ assetId: 'asset-1', personId: 'person-1' }]);
+    });
+  });
+
+  describe('handleRemoveFace', () => {
+    it('should skip if search is disabled', () => {
+      const sut = makeSut('false');
+      sut.handleRemoveFace({ assetId: 'asset-1', personId: 'person-1' });
+    });
+
+    it('should remove the face', () => {
+      sut.handleRemoveFace({ assetId: 'asset-1', personId: 'person-1' });
+    });
+  });
+
  describe('flush', () => {
    it('should flush queued album updates', async () => {
      albumMock.getByIds.mockResolvedValue([albumStub.empty]);
@@ -1,4 +1,4 @@
-import { AlbumEntity, AssetEntity } from '@app/infra/entities';
+import { AlbumEntity, AssetEntity, AssetFaceEntity } from '@app/infra/entities';
 import { BadRequestException, Inject, Injectable, Logger } from '@nestjs/common';
 import { ConfigService } from '@nestjs/config';
 import { mapAlbum } from '../album';
@@ -7,12 +7,14 @@ import { mapAsset } from '../asset';
 import { IAssetRepository } from '../asset/asset.repository';
 import { AuthUserDto } from '../auth';
 import { MACHINE_LEARNING_ENABLED } from '../domain.constant';
-import { IBulkEntityJob, IJobRepository, JobName } from '../job';
+import { AssetFaceId, IFaceRepository } from '../facial-recognition';
+import { IAssetFaceJob, IBulkEntityJob, IJobRepository, JobName } from '../job';
 import { IMachineLearningRepository } from '../smart-info';
 import { SearchDto } from './dto';
 import { SearchConfigResponseDto, SearchResponseDto } from './response-dto';
 import {
  ISearchRepository,
+  OwnedFaceEntity,
  SearchCollection,
  SearchExploreItem,
  SearchResult,
@@ -40,9 +42,15 @@ export class SearchService {
    delete: new Set(),
  };

+  private faceQueue: SyncQueue = {
+    upsert: new Set(),
+    delete: new Set(),
+  };
+
  constructor(
    @Inject(IAlbumRepository) private albumRepository: IAlbumRepository,
    @Inject(IAssetRepository) private assetRepository: IAssetRepository,
+    @Inject(IFaceRepository) private faceRepository: IFaceRepository,
    @Inject(IJobRepository) private jobRepository: IJobRepository,
    @Inject(IMachineLearningRepository) private machineLearning: IMachineLearningRepository,
    @Inject(ISearchRepository) private searchRepository: ISearchRepository,
@@ -88,6 +96,10 @@ export class SearchService {
      this.logger.debug('Queueing job to re-index all albums');
      await this.jobRepository.queue({ name: JobName.SEARCH_INDEX_ALBUMS });
    }
+    if (migrationStatus[SearchCollection.FACES]) {
+      this.logger.debug('Queueing job to re-index all faces');
+      await this.jobRepository.queue({ name: JobName.SEARCH_INDEX_FACES });
+    }
  }

  async getExploreData(authUser: AuthUserDto): Promise<SearchExploreItem<AssetEntity>[]> {
@@ -159,6 +171,29 @@ export class SearchService {
    }
  }

+  async handleIndexFaces() {
+    if (!this.enabled) {
+      return;
+    }
+
+    try {
+      // TODO: do this in batches based on searchIndexVersion
+      const faces = this.patchFaces(await this.faceRepository.getAll());
+      this.logger.log(`Indexing ${faces.length} faces`);
+
+      const chunkSize = 1000;
+      for (let i = 0; i < faces.length; i += chunkSize) {
+        await this.searchRepository.importFaces(faces.slice(i, i + chunkSize), false);
+      }
+
+      await this.searchRepository.importFaces([], true);
+
+      this.logger.debug('Finished re-indexing all faces');
+    } catch (error: any) {
+      this.logger.error(`Unable to index all faces`, error?.stack);
+    }
+  }
+
  handleIndexAlbum({ ids }: IBulkEntityJob) {
    if (!this.enabled) {
      return;
@@ -179,6 +214,15 @@ export class SearchService {
    }
  }

+  async handleIndexFace({ assetId, personId }: IAssetFaceJob) {
+    if (!this.enabled) {
+      return;
+    }
+
+    // immediately push to typesense
+    await this.searchRepository.importFaces(await this.idsToFaces([{ assetId, personId }]), false);
+  }
+
  handleRemoveAlbum({ ids }: IBulkEntityJob) {
    if (!this.enabled) {
      return;
@@ -199,6 +243,14 @@ export class SearchService {
    }
  }

+  handleRemoveFace({ assetId, personId }: IAssetFaceJob) {
+    if (!this.enabled) {
+      return;
+    }
+
+    this.faceQueue.delete.add(this.asKey({ assetId, personId }));
+  }
+
  private async flush() {
    if (this.albumQueue.upsert.size > 0) {
      const ids = [...this.albumQueue.upsert.keys()];
@@ -229,6 +281,21 @@ export class SearchService {
      await this.searchRepository.deleteAssets(ids);
      this.assetQueue.delete.clear();
    }
+
+    if (this.faceQueue.upsert.size > 0) {
+      const ids = [...this.faceQueue.upsert.keys()].map((key) => this.asParts(key));
+      const items = await this.idsToFaces(ids);
+      this.logger.debug(`Flushing ${items.length} face upserts`);
+      await this.searchRepository.importFaces(items, false);
+      this.faceQueue.upsert.clear();
+    }
+
+    if (this.faceQueue.delete.size > 0) {
+      const ids = [...this.faceQueue.delete.keys()];
+      this.logger.debug(`Flushing ${ids.length} face deletes`);
+      await this.searchRepository.deleteFaces(ids);
+      this.faceQueue.delete.clear();
+    }
  }

  private assertEnabled() {
@@ -247,6 +314,10 @@ export class SearchService {
    return this.patchAssets(entities.filter((entity) => entity.isVisible));
  }

+  private async idsToFaces(ids: AssetFaceId[]): Promise<OwnedFaceEntity[]> {
+    return this.patchFaces(await this.faceRepository.getByIds(ids));
+  }
+
  private patchAssets(assets: AssetEntity[]): AssetEntity[] {
    return assets;
  }
@@ -254,4 +325,23 @@ export class SearchService {
  private patchAlbums(albums: AlbumEntity[]): AlbumEntity[] {
    return albums.map((entity) => ({ ...entity, assets: [] }));
  }
+
+  private patchFaces(faces: AssetFaceEntity[]): OwnedFaceEntity[] {
+    return faces.map((face) => ({
+      id: this.asKey(face),
+      ownerId: face.asset.ownerId,
+      assetId: face.assetId,
+      personId: face.personId,
+      embedding: face.embedding,
+    }));
+  }
+
+  private asKey(face: AssetFaceId): string {
+    return `${face.assetId}|${face.personId}`;
+  }
+
+  private asParts(key: string): AssetFaceId {
+    const [assetId, personId] = key.split('|');
+    return { assetId, personId };
+  }
 }