feat: facial recognition (#2180)

This commit is contained in:
Jason Rasmussen
2023-05-17 13:07:17 -04:00
committed by GitHub
parent 115a47d4c6
commit 93863b0629
107 changed files with 3943 additions and 133 deletions

View File

@@ -18,6 +18,7 @@ export enum WithoutProperty {
EXIF = 'exif',
CLIP_ENCODING = 'clip-embedding',
OBJECT_TAGS = 'object-tags',
FACES = 'faces',
}
export const IAssetRepository = 'IAssetRepository';

View File

@@ -1,8 +1,9 @@
import { AssetEntity, AssetType } from '@app/infra/entities';
import { ApiProperty } from '@nestjs/swagger';
import { mapFace, PersonResponseDto } from '../../person';
import { mapTag, TagResponseDto } from '../../tag';
import { ExifResponseDto, mapExif } from './exif-response.dto';
import { SmartInfoResponseDto, mapSmartInfo } from './smart-info-response.dto';
import { mapSmartInfo, SmartInfoResponseDto } from './smart-info-response.dto';
export class AssetResponseDto {
id!: string;
@@ -28,6 +29,7 @@ export class AssetResponseDto {
smartInfo?: SmartInfoResponseDto;
livePhotoVideoId?: string | null;
tags?: TagResponseDto[];
people?: PersonResponseDto[];
}
export function mapAsset(entity: AssetEntity): AssetResponseDto {
@@ -53,6 +55,7 @@ export function mapAsset(entity: AssetEntity): AssetResponseDto {
smartInfo: entity.smartInfo ? mapSmartInfo(entity.smartInfo) : undefined,
livePhotoVideoId: entity.livePhotoVideoId,
tags: entity.tags?.map(mapTag),
people: entity.faces?.map(mapFace),
};
}
@@ -79,5 +82,6 @@ export function mapAssetWithoutExif(entity: AssetEntity): AssetResponseDto {
smartInfo: entity.smartInfo ? mapSmartInfo(entity.smartInfo) : undefined,
livePhotoVideoId: entity.livePhotoVideoId,
tags: entity.tags?.map(mapTag),
people: entity.faces?.map(mapFace),
};
}

View File

@@ -3,27 +3,31 @@ import { AlbumService } from './album';
import { APIKeyService } from './api-key';
import { AssetService } from './asset';
import { AuthService } from './auth';
import { FacialRecognitionService } from './facial-recognition';
import { JobService } from './job';
import { MediaService } from './media';
import { OAuthService } from './oauth';
import { PartnerService } from './partner';
import { PersonService } from './person';
import { SearchService } from './search';
import { ServerInfoService } from './server-info';
import { ShareService } from './share';
import { SmartInfoService } from './smart-info';
import { StorageService } from './storage';
import { StorageTemplateService } from './storage-template';
import { UserService } from './user';
import { INITIAL_SYSTEM_CONFIG, SystemConfigService } from './system-config';
import { UserService } from './user';
const providers: Provider[] = [
AlbumService,
APIKeyService,
AssetService,
AuthService,
FacialRecognitionService,
JobService,
MediaService,
OAuthService,
PersonService,
PartnerService,
SearchService,
ServerInfoService,

View File

@@ -0,0 +1,14 @@
import { AssetFaceEntity } from '@app/infra/entities';
export const IFaceRepository = 'IFaceRepository';
export interface AssetFaceId {
assetId: string;
personId: string;
}
export interface IFaceRepository {
getAll(): Promise<AssetFaceEntity[]>;
getByIds(ids: AssetFaceId[]): Promise<AssetFaceEntity[]>;
create(entity: Partial<AssetFaceEntity>): Promise<AssetFaceEntity>;
}

View File

@@ -0,0 +1,320 @@
import {
assetEntityStub,
faceStub,
newAssetRepositoryMock,
newFaceRepositoryMock,
newJobRepositoryMock,
newMachineLearningRepositoryMock,
newMediaRepositoryMock,
newPersonRepositoryMock,
newSearchRepositoryMock,
newStorageRepositoryMock,
personStub,
} from '../../test';
import { IAssetRepository, WithoutProperty } from '../asset';
import { IJobRepository, JobName } from '../job';
import { IMediaRepository } from '../media';
import { IPersonRepository } from '../person';
import { ISearchRepository } from '../search';
import { IMachineLearningRepository } from '../smart-info';
import { IStorageRepository } from '../storage';
import { IFaceRepository } from './face.repository';
import { FacialRecognitionService } from './facial-recognition.services';
const croppedFace = Buffer.from('Cropped Face');
const face = {
start: {
assetId: 'asset-1',
personId: 'person-1',
boundingBox: {
x1: 5,
y1: 5,
x2: 505,
y2: 505,
},
imageHeight: 1000,
imageWidth: 1000,
},
middle: {
assetId: 'asset-1',
personId: 'person-1',
boundingBox: {
x1: 100,
y1: 100,
x2: 200,
y2: 200,
},
imageHeight: 500,
imageWidth: 400,
embedding: [1, 2, 3, 4],
score: 0.2,
},
end: {
assetId: 'asset-1',
personId: 'person-1',
boundingBox: {
x1: 300,
y1: 300,
x2: 495,
y2: 495,
},
imageHeight: 500,
imageWidth: 500,
},
};
const faceSearch = {
noMatch: {
total: 0,
count: 0,
page: 1,
items: [],
distances: [],
facets: [],
},
oneMatch: {
total: 1,
count: 1,
page: 1,
items: [faceStub.face1],
distances: [0.1],
facets: [],
},
oneRemoteMatch: {
total: 1,
count: 1,
page: 1,
items: [faceStub.face1],
distances: [0.8],
facets: [],
},
};
describe(FacialRecognitionService.name, () => {
let sut: FacialRecognitionService;
let assetMock: jest.Mocked<IAssetRepository>;
let faceMock: jest.Mocked<IFaceRepository>;
let jobMock: jest.Mocked<IJobRepository>;
let machineLearningMock: jest.Mocked<IMachineLearningRepository>;
let mediaMock: jest.Mocked<IMediaRepository>;
let personMock: jest.Mocked<IPersonRepository>;
let searchMock: jest.Mocked<ISearchRepository>;
let storageMock: jest.Mocked<IStorageRepository>;
beforeEach(async () => {
assetMock = newAssetRepositoryMock();
faceMock = newFaceRepositoryMock();
jobMock = newJobRepositoryMock();
machineLearningMock = newMachineLearningRepositoryMock();
mediaMock = newMediaRepositoryMock();
personMock = newPersonRepositoryMock();
searchMock = newSearchRepositoryMock();
storageMock = newStorageRepositoryMock();
mediaMock.crop.mockResolvedValue(croppedFace);
sut = new FacialRecognitionService(
assetMock,
faceMock,
jobMock,
machineLearningMock,
mediaMock,
personMock,
searchMock,
storageMock,
);
});
it('should be defined', () => {
expect(sut).toBeDefined();
});
describe('handleQueueRecognizeFaces', () => {
it('should queue missing assets', async () => {
assetMock.getWithout.mockResolvedValue([assetEntityStub.image]);
await sut.handleQueueRecognizeFaces({});
expect(assetMock.getWithout).toHaveBeenCalledWith(WithoutProperty.FACES);
expect(jobMock.queue).toHaveBeenCalledWith({
name: JobName.RECOGNIZE_FACES,
data: { asset: assetEntityStub.image },
});
});
it('should queue all assets', async () => {
assetMock.getAll.mockResolvedValue([assetEntityStub.image]);
personMock.deleteAll.mockResolvedValue(5);
searchMock.deleteAllFaces.mockResolvedValue(100);
await sut.handleQueueRecognizeFaces({ force: true });
expect(assetMock.getAll).toHaveBeenCalled();
expect(jobMock.queue).toHaveBeenCalledWith({
name: JobName.RECOGNIZE_FACES,
data: { asset: assetEntityStub.image },
});
});
it('should log an error', async () => {
assetMock.getWithout.mockRejectedValue(new Error('Database unavailable'));
await sut.handleQueueRecognizeFaces({});
});
});
describe('handleRecognizeFaces', () => {
it('should skip when no resize path', async () => {
await sut.handleRecognizeFaces({ asset: assetEntityStub.noResizePath });
expect(machineLearningMock.detectFaces).not.toHaveBeenCalled();
});
it('should handle no results', async () => {
machineLearningMock.detectFaces.mockResolvedValue([]);
await sut.handleRecognizeFaces({ asset: assetEntityStub.image });
expect(machineLearningMock.detectFaces).toHaveBeenCalledWith({
thumbnailPath: assetEntityStub.image.resizePath,
});
expect(faceMock.create).not.toHaveBeenCalled();
expect(jobMock.queue).not.toHaveBeenCalled();
});
it('should match existing people', async () => {
machineLearningMock.detectFaces.mockResolvedValue([face.middle]);
searchMock.searchFaces.mockResolvedValue(faceSearch.oneMatch);
await sut.handleRecognizeFaces({ asset: assetEntityStub.image });
expect(faceMock.create).toHaveBeenCalledWith({
personId: 'person-1',
assetId: 'asset-id',
embedding: [1, 2, 3, 4],
});
expect(jobMock.queue.mock.calls).toEqual([
[{ name: JobName.SEARCH_INDEX_FACE, data: { personId: 'person-1', assetId: 'asset-id' } }],
[{ name: JobName.SEARCH_INDEX_ASSET, data: { ids: ['asset-id'] } }],
]);
});
it('should create a new person', async () => {
machineLearningMock.detectFaces.mockResolvedValue([face.middle]);
searchMock.searchFaces.mockResolvedValue(faceSearch.oneRemoteMatch);
personMock.create.mockResolvedValue(personStub.noName);
await sut.handleRecognizeFaces({ asset: assetEntityStub.image });
expect(personMock.create).toHaveBeenCalledWith({ ownerId: assetEntityStub.image.ownerId });
expect(faceMock.create).toHaveBeenCalledWith({
personId: 'person-1',
assetId: 'asset-id',
embedding: [1, 2, 3, 4],
});
expect(jobMock.queue.mock.calls).toEqual([
[
{
name: JobName.GENERATE_FACE_THUMBNAIL,
data: {
assetId: 'asset-1',
personId: 'person-1',
boundingBox: {
x1: 100,
y1: 100,
x2: 200,
y2: 200,
},
imageHeight: 500,
imageWidth: 400,
score: 0.2,
},
},
],
[{ name: JobName.SEARCH_INDEX_FACE, data: { personId: 'person-1', assetId: 'asset-id' } }],
[{ name: JobName.SEARCH_INDEX_ASSET, data: { ids: ['asset-id'] } }],
]);
});
it('should log an error', async () => {
machineLearningMock.detectFaces.mockRejectedValue(new Error('machine learning unavailable'));
await sut.handleRecognizeFaces({ asset: assetEntityStub.image });
});
});
describe('handleGenerateFaceThumbnail', () => {
it('should skip an asset not found', async () => {
assetMock.getByIds.mockResolvedValue([]);
await sut.handleGenerateFaceThumbnail(face.middle);
expect(mediaMock.crop).not.toHaveBeenCalled();
});
it('should skip an asset without a thumbnail', async () => {
assetMock.getByIds.mockResolvedValue([assetEntityStub.noResizePath]);
await sut.handleGenerateFaceThumbnail(face.middle);
expect(mediaMock.crop).not.toHaveBeenCalled();
});
it('should generate a thumbnail', async () => {
assetMock.getByIds.mockResolvedValue([assetEntityStub.image]);
await sut.handleGenerateFaceThumbnail(face.middle);
expect(assetMock.getByIds).toHaveBeenCalledWith(['asset-1']);
expect(storageMock.mkdirSync).toHaveBeenCalledWith('upload/thumbs/user-id');
expect(mediaMock.crop).toHaveBeenCalledWith('/uploads/user-id/thumbs/path.ext', {
left: 95,
top: 95,
width: 110,
height: 110,
});
expect(mediaMock.resize).toHaveBeenCalledWith(croppedFace, 'upload/thumbs/user-id/person-1.jpeg', {
format: 'jpeg',
size: 250,
});
expect(personMock.update).toHaveBeenCalledWith({
id: 'person-1',
thumbnailPath: 'upload/thumbs/user-id/person-1.jpeg',
});
});
it('should generate a thumbnail without going negative', async () => {
assetMock.getByIds.mockResolvedValue([assetEntityStub.image]);
await sut.handleGenerateFaceThumbnail(face.start);
expect(mediaMock.crop).toHaveBeenCalledWith('/uploads/user-id/thumbs/path.ext', {
left: 0,
top: 0,
width: 510,
height: 510,
});
expect(mediaMock.resize).toHaveBeenCalledWith(croppedFace, 'upload/thumbs/user-id/person-1.jpeg', {
format: 'jpeg',
size: 250,
});
});
it('should generate a thumbnail without overflowing', async () => {
assetMock.getByIds.mockResolvedValue([assetEntityStub.image]);
await sut.handleGenerateFaceThumbnail(face.end);
expect(mediaMock.crop).toHaveBeenCalledWith('/uploads/user-id/thumbs/path.ext', {
left: 297,
top: 297,
width: 202,
height: 202,
});
expect(mediaMock.resize).toHaveBeenCalledWith(croppedFace, 'upload/thumbs/user-id/person-1.jpeg', {
format: 'jpeg',
size: 250,
});
});
it('should log an error', async () => {
assetMock.getByIds.mockRejectedValue(new Error('Database unavailable'));
await sut.handleGenerateFaceThumbnail(face.middle);
});
});
});

View File

@@ -0,0 +1,144 @@
import { Inject, Logger } from '@nestjs/common';
import { join } from 'path';
import { IAssetRepository, WithoutProperty } from '../asset';
import { MACHINE_LEARNING_ENABLED } from '../domain.constant';
import { IAssetJob, IBaseJob, IFaceThumbnailJob, IJobRepository, JobName } from '../job';
import { CropOptions, FACE_THUMBNAIL_SIZE, IMediaRepository } from '../media';
import { IPersonRepository } from '../person/person.repository';
import { ISearchRepository } from '../search/search.repository';
import { IMachineLearningRepository } from '../smart-info';
import { IStorageRepository, StorageCore, StorageFolder } from '../storage';
import { AssetFaceId, IFaceRepository } from './face.repository';
export class FacialRecognitionService {
private logger = new Logger(FacialRecognitionService.name);
private storageCore = new StorageCore();
constructor(
@Inject(IAssetRepository) private assetRepository: IAssetRepository,
@Inject(IFaceRepository) private faceRepository: IFaceRepository,
@Inject(IJobRepository) private jobRepository: IJobRepository,
@Inject(IMachineLearningRepository) private machineLearning: IMachineLearningRepository,
@Inject(IMediaRepository) private mediaRepository: IMediaRepository,
@Inject(IPersonRepository) private personRepository: IPersonRepository,
@Inject(ISearchRepository) private searchRepository: ISearchRepository,
@Inject(IStorageRepository) private storageRepository: IStorageRepository,
) {}
async handleQueueRecognizeFaces({ force }: IBaseJob) {
try {
const assets = force
? await this.assetRepository.getAll()
: await this.assetRepository.getWithout(WithoutProperty.FACES);
if (force) {
const people = await this.personRepository.deleteAll();
const faces = await this.searchRepository.deleteAllFaces();
this.logger.debug(`Deleted ${people} people and ${faces} faces`);
}
for (const asset of assets) {
await this.jobRepository.queue({ name: JobName.RECOGNIZE_FACES, data: { asset } });
}
} catch (error: any) {
this.logger.error(`Unable to queue recognize faces`, error?.stack);
}
}
async handleRecognizeFaces(data: IAssetJob) {
const { asset } = data;
if (!MACHINE_LEARNING_ENABLED || !asset.resizePath) {
return;
}
try {
const faces = await this.machineLearning.detectFaces({ thumbnailPath: asset.resizePath });
this.logger.debug(`${faces.length} faces detected in ${asset.resizePath}`);
this.logger.verbose(faces.map((face) => ({ ...face, embedding: `float[${face.embedding.length}]` })));
for (const { embedding, ...rest } of faces) {
const faceSearchResult = await this.searchRepository.searchFaces(embedding, { ownerId: asset.ownerId });
let personId: string | null = null;
// try to find a matching face and link to the associated person
// The closer to 0, the better the match. Range is from 0 to 2
if (faceSearchResult.total && faceSearchResult.distances[0] < 0.6) {
this.logger.verbose(`Match face with distance ${faceSearchResult.distances[0]}`);
personId = faceSearchResult.items[0].personId;
}
if (!personId) {
this.logger.debug('No matches, creating a new person.');
const person = await this.personRepository.create({ ownerId: asset.ownerId });
personId = person.id;
await this.jobRepository.queue({
name: JobName.GENERATE_FACE_THUMBNAIL,
data: { assetId: asset.id, personId, ...rest },
});
}
const faceId: AssetFaceId = { assetId: asset.id, personId };
await this.faceRepository.create({ ...faceId, embedding });
await this.jobRepository.queue({ name: JobName.SEARCH_INDEX_FACE, data: faceId });
await this.jobRepository.queue({ name: JobName.SEARCH_INDEX_ASSET, data: { ids: [asset.id] } });
}
// queue all faces for asset
} catch (error: any) {
this.logger.error(`Unable run facial recognition pipeline: ${asset.id}`, error?.stack);
}
}
async handleGenerateFaceThumbnail(data: IFaceThumbnailJob) {
const { assetId, personId, boundingBox, imageWidth, imageHeight } = data;
try {
const [asset] = await this.assetRepository.getByIds([assetId]);
if (!asset || !asset.resizePath) {
this.logger.warn(`Asset not found for facial cropping: ${assetId}`);
return null;
}
this.logger.verbose(`Cropping face for person: ${personId}`);
const outputFolder = this.storageCore.getFolderLocation(StorageFolder.THUMBNAILS, asset.ownerId);
const output = join(outputFolder, `${personId}.jpeg`);
this.storageRepository.mkdirSync(outputFolder);
const { x1, y1, x2, y2 } = boundingBox;
const halfWidth = (x2 - x1) / 2;
const halfHeight = (y2 - y1) / 2;
const middleX = Math.round(x1 + halfWidth);
const middleY = Math.round(y1 + halfHeight);
// zoom out 10%
const targetHalfSize = Math.floor(Math.max(halfWidth, halfHeight) * 1.1);
// get the longest distance from the center of the image without overflowing
const newHalfSize = Math.min(
middleX - Math.max(0, middleX - targetHalfSize),
middleY - Math.max(0, middleY - targetHalfSize),
Math.min(imageWidth - 1, middleX + targetHalfSize) - middleX,
Math.min(imageHeight - 1, middleY + targetHalfSize) - middleY,
);
const cropOptions: CropOptions = {
left: middleX - newHalfSize,
top: middleY - newHalfSize,
width: newHalfSize * 2,
height: newHalfSize * 2,
};
const croppedOutput = await this.mediaRepository.crop(asset.resizePath, cropOptions);
await this.mediaRepository.resize(croppedOutput, output, { size: FACE_THUMBNAIL_SIZE, format: 'jpeg' });
await this.personRepository.update({ id: personId, thumbnailPath: output });
} catch (error: Error | any) {
this.logger.error(`Failed to crop face for asset: ${assetId}, person: ${personId} - ${error}`, error.stack);
}
}
}

View File

@@ -0,0 +1,2 @@
export * from './facial-recognition.services';
export * from './face.repository';

View File

@@ -8,10 +8,12 @@ export * from './domain.config';
export * from './domain.constant';
export * from './domain.module';
export * from './domain.util';
export * from './facial-recognition';
export * from './job';
export * from './media';
export * from './metadata';
export * from './oauth';
export * from './person';
export * from './search';
export * from './server-info';
export * from './partner';

View File

@@ -3,6 +3,7 @@ export enum QueueName {
METADATA_EXTRACTION = 'metadata-extraction-queue',
VIDEO_CONVERSION = 'video-conversion-queue',
OBJECT_TAGGING = 'object-tagging-queue',
RECOGNIZE_FACES = 'recognize-faces-queue',
CLIP_ENCODING = 'clip-encoding-queue',
BACKGROUND_TASK = 'background-task-queue',
STORAGE_TEMPLATE_MIGRATION = 'storage-template-migration-queue',
@@ -48,16 +49,25 @@ export enum JobName {
DETECT_OBJECTS = 'detect-objects',
CLASSIFY_IMAGE = 'classify-image',
// facial recognition
QUEUE_RECOGNIZE_FACES = 'queue-recognize-faces',
RECOGNIZE_FACES = 'recognize-faces',
GENERATE_FACE_THUMBNAIL = 'generate-face-thumbnail',
PERSON_CLEANUP = 'person-cleanup',
// cleanup
DELETE_FILES = 'delete-files',
// search
SEARCH_INDEX_ASSETS = 'search-index-assets',
SEARCH_INDEX_ASSET = 'search-index-asset',
SEARCH_INDEX_FACE = 'search-index-face',
SEARCH_INDEX_FACES = 'search-index-faces',
SEARCH_INDEX_ALBUMS = 'search-index-albums',
SEARCH_INDEX_ALBUM = 'search-index-album',
SEARCH_REMOVE_ALBUM = 'search-remove-album',
SEARCH_REMOVE_ASSET = 'search-remove-asset',
SEARCH_REMOVE_FACE = 'search-remove-face',
// clip
QUEUE_ENCODE_CLIP = 'queue-clip-encode',

View File

@@ -1,4 +1,5 @@
import { AlbumEntity, AssetEntity, UserEntity } from '@app/infra/entities';
import { BoundingBox } from '../smart-info';
export interface IBaseJob {
force?: boolean;
@@ -12,6 +13,19 @@ export interface IAssetJob extends IBaseJob {
asset: AssetEntity;
}
export interface IAssetFaceJob extends IBaseJob {
assetId: string;
personId: string;
}
export interface IFaceThumbnailJob extends IAssetFaceJob {
imageWidth: number;
imageHeight: number;
boundingBox: BoundingBox;
assetId: string;
personId: string;
}
export interface IBulkEntityJob extends IBaseJob {
ids: string[];
}

View File

@@ -1,10 +1,12 @@
import { JobName, QueueName } from './job.constants';
import {
IAssetFaceJob,
IAssetJob,
IAssetUploadedJob,
IBaseJob,
IBulkEntityJob,
IDeleteFilesJob,
IFaceThumbnailJob,
IUserDeletionJob,
} from './job.interface';
@@ -54,6 +56,11 @@ export type JobItem =
| { name: JobName.DETECT_OBJECTS; data: IAssetJob }
| { name: JobName.CLASSIFY_IMAGE; data: IAssetJob }
// Recognize Faces
| { name: JobName.QUEUE_RECOGNIZE_FACES; data: IBaseJob }
| { name: JobName.RECOGNIZE_FACES; data: IAssetJob }
| { name: JobName.GENERATE_FACE_THUMBNAIL; data: IFaceThumbnailJob }
// Clip Embedding
| { name: JobName.QUEUE_ENCODE_CLIP; data: IBaseJob }
| { name: JobName.ENCODE_CLIP; data: IAssetJob }
@@ -61,13 +68,19 @@ export type JobItem =
// Filesystem
| { name: JobName.DELETE_FILES; data: IDeleteFilesJob }
// Asset Deletion
| { name: JobName.PERSON_CLEANUP }
// Search
| { name: JobName.SEARCH_INDEX_ASSETS }
| { name: JobName.SEARCH_INDEX_ASSET; data: IBulkEntityJob }
| { name: JobName.SEARCH_INDEX_FACES }
| { name: JobName.SEARCH_INDEX_FACE; data: IAssetFaceJob }
| { name: JobName.SEARCH_INDEX_ALBUMS }
| { name: JobName.SEARCH_INDEX_ALBUM; data: IBulkEntityJob }
| { name: JobName.SEARCH_REMOVE_ASSET; data: IBulkEntityJob }
| { name: JobName.SEARCH_REMOVE_ALBUM; data: IBulkEntityJob };
| { name: JobName.SEARCH_REMOVE_ALBUM; data: IBulkEntityJob }
| { name: JobName.SEARCH_REMOVE_FACE; data: IAssetFaceJob };
export const IJobRepository = 'IJobRepository';

View File

@@ -15,6 +15,17 @@ describe(JobService.name, () => {
expect(sut).toBeDefined();
});
describe('handleNightlyJobs', () => {
it('should run the scheduled jobs', async () => {
await sut.handleNightlyJobs();
expect(jobMock.queue.mock.calls).toEqual([
[{ name: JobName.USER_DELETE_CHECK }],
[{ name: JobName.PERSON_CLEANUP }],
]);
});
});
describe('getAllJobStatus', () => {
it('should get all job statuses', async () => {
jobMock.getJobCounts.mockResolvedValue({
@@ -54,6 +65,7 @@ describe(JobService.name, () => {
'storage-template-migration-queue': expectedJobStatus,
'thumbnail-generation-queue': expectedJobStatus,
'video-conversion-queue': expectedJobStatus,
'recognize-faces-queue': expectedJobStatus,
});
});
});

View File

@@ -11,6 +11,11 @@ export class JobService {
constructor(@Inject(IJobRepository) private jobRepository: IJobRepository) {}
async handleNightlyJobs() {
await this.jobRepository.queue({ name: JobName.USER_DELETE_CHECK });
await this.jobRepository.queue({ name: JobName.PERSON_CLEANUP });
}
handleCommand(queueName: QueueName, dto: JobCommandDto): Promise<void> {
this.logger.debug(`Handling command: queue=${queueName},force=${dto.force}`);
@@ -73,6 +78,9 @@ export class JobService {
case QueueName.THUMBNAIL_GENERATION:
return this.jobRepository.queue({ name: JobName.QUEUE_GENERATE_THUMBNAILS, data: { force } });
case QueueName.RECOGNIZE_FACES:
return this.jobRepository.queue({ name: JobName.QUEUE_RECOGNIZE_FACES, data: { force } });
default:
throw new BadRequestException(`Invalid job name: ${name}`);
}

View File

@@ -53,4 +53,7 @@ export class AllJobStatusResponseDto implements Record<QueueName, JobStatusDto>
@ApiProperty({ type: JobStatusDto })
[QueueName.SEARCH]!: JobStatusDto;
@ApiProperty({ type: JobStatusDto })
[QueueName.RECOGNIZE_FACES]!: JobStatusDto;
}

View File

@@ -1,2 +1,3 @@
export * from './media.constant';
export * from './media.repository';
export * from './media.service';

View File

@@ -0,0 +1,3 @@
export const JPEG_THUMBNAIL_SIZE = 1440;
export const WEBP_THUMBNAIL_SIZE = 250;
export const FACE_THUMBNAIL_SIZE = 250;

View File

@@ -31,10 +31,18 @@ export interface VideoInfo {
audioStreams: AudioStreamInfo[];
}
export interface CropOptions {
top: number;
left: number;
width: number;
height: number;
}
export interface IMediaRepository {
// image
extractThumbnailFromExif(input: string, output: string): Promise<void>;
resize(input: string, output: string, options: ResizeOptions): Promise<void>;
resize(input: string | Buffer, output: string, options: ResizeOptions): Promise<void>;
crop(input: string, options: CropOptions): Promise<Buffer>;
// video
extractVideoThumbnail(input: string, output: string, size: number): Promise<void>;

View File

@@ -34,6 +34,7 @@ describe(MediaService.name, () => {
jobMock = newJobRepositoryMock();
mediaMock = newMediaRepositoryMock();
storageMock = newStorageRepositoryMock();
sut = new MediaService(assetMock, communicationMock, jobMock, mediaMock, storageMock, configMock);
});

View File

@@ -7,6 +7,7 @@ import { IAssetJob, IBaseJob, IJobRepository, JobName } from '../job';
import { IStorageRepository, StorageCore, StorageFolder } from '../storage';
import { ISystemConfigRepository, SystemConfigFFmpegDto } from '../system-config';
import { SystemConfigCore } from '../system-config/system-config.core';
import { JPEG_THUMBNAIL_SIZE, WEBP_THUMBNAIL_SIZE } from './media.constant';
import { AudioStreamInfo, IMediaRepository, VideoStreamInfo } from './media.repository';
@Injectable()
@@ -57,11 +58,10 @@ export class MediaService {
this.storageRepository.mkdirSync(resizePath);
const jpegThumbnailPath = join(resizePath, `${asset.id}.jpeg`);
const thumbnailDimension = 1440;
if (asset.type == AssetType.IMAGE) {
try {
await this.mediaRepository.resize(asset.originalPath, jpegThumbnailPath, {
size: thumbnailDimension,
size: JPEG_THUMBNAIL_SIZE,
format: 'jpeg',
});
} catch (error) {
@@ -74,7 +74,7 @@ export class MediaService {
if (asset.type == AssetType.VIDEO) {
this.logger.log('Start Generating Video Thumbnail');
await this.mediaRepository.extractVideoThumbnail(asset.originalPath, jpegThumbnailPath, thumbnailDimension);
await this.mediaRepository.extractVideoThumbnail(asset.originalPath, jpegThumbnailPath, JPEG_THUMBNAIL_SIZE);
this.logger.log(`Generating Video Thumbnail Success ${asset.id}`);
}
@@ -86,6 +86,7 @@ export class MediaService {
await this.jobRepository.queue({ name: JobName.CLASSIFY_IMAGE, data: { asset } });
await this.jobRepository.queue({ name: JobName.DETECT_OBJECTS, data: { asset } });
await this.jobRepository.queue({ name: JobName.ENCODE_CLIP, data: { asset } });
await this.jobRepository.queue({ name: JobName.RECOGNIZE_FACES, data: { asset } });
this.communicationRepository.send(CommunicationEvent.UPLOAD_SUCCESS, asset.ownerId, mapAsset(asset));
} catch (error: any) {
@@ -103,7 +104,7 @@ export class MediaService {
const webpPath = asset.resizePath.replace('jpeg', 'webp');
try {
await this.mediaRepository.resize(asset.resizePath, webpPath, { size: 250, format: 'webp' });
await this.mediaRepository.resize(asset.resizePath, webpPath, { size: WEBP_THUMBNAIL_SIZE, format: 'webp' });
await this.assetRepository.save({ id: asset.id, webpPath: webpPath });
} catch (error: any) {
this.logger.error(`Failed to generate webp thumbnail for asset: ${asset.id}`, error.stack);

View File

@@ -0,0 +1 @@
export * from './person-update.dto';

View File

@@ -0,0 +1,7 @@
import { IsNotEmpty, IsString } from 'class-validator';
export class PersonUpdateDto {
@IsNotEmpty()
@IsString()
name!: string;
}

View File

@@ -0,0 +1,4 @@
export * from './dto';
export * from './person.repository';
export * from './person.service';
export * from './response-dto';

View File

@@ -0,0 +1,19 @@
import { AssetEntity, PersonEntity } from '@app/infra/entities';
export const IPersonRepository = 'IPersonRepository';
export interface PersonSearchOptions {
minimumFaceCount: number;
}
export interface IPersonRepository {
getAll(userId: string, options: PersonSearchOptions): Promise<PersonEntity[]>;
getAllWithoutFaces(): Promise<PersonEntity[]>;
getById(userId: string, personId: string): Promise<PersonEntity | null>;
getAssets(userId: string, id: string): Promise<AssetEntity[]>;
create(entity: Partial<PersonEntity>): Promise<PersonEntity>;
update(entity: Partial<PersonEntity>): Promise<PersonEntity>;
delete(entity: PersonEntity): Promise<PersonEntity | null>;
deleteAll(): Promise<number>;
}

View File

@@ -0,0 +1,135 @@
import { BadRequestException, NotFoundException } from '@nestjs/common';
import { IJobRepository, JobName } from '..';
import {
assetEntityStub,
authStub,
newJobRepositoryMock,
newPersonRepositoryMock,
newStorageRepositoryMock,
personStub,
} from '../../test';
import { IStorageRepository } from '../storage';
import { IPersonRepository } from './person.repository';
import { PersonService } from './person.service';
import { PersonResponseDto } from './response-dto';
const responseDto: PersonResponseDto = {
id: 'person-1',
name: 'Person 1',
thumbnailPath: '/path/to/thumbnail',
};
describe(PersonService.name, () => {
let sut: PersonService;
let personMock: jest.Mocked<IPersonRepository>;
let storageMock: jest.Mocked<IStorageRepository>;
let jobMock: jest.Mocked<IJobRepository>;
beforeEach(async () => {
personMock = newPersonRepositoryMock();
storageMock = newStorageRepositoryMock();
jobMock = newJobRepositoryMock();
sut = new PersonService(personMock, storageMock, jobMock);
});
it('should be defined', () => {
expect(sut).toBeDefined();
});
describe('getAll', () => {
it('should get all people with thumbnails', async () => {
personMock.getAll.mockResolvedValue([personStub.withName, personStub.noThumbnail]);
await expect(sut.getAll(authStub.admin)).resolves.toEqual([responseDto]);
expect(personMock.getAll).toHaveBeenCalledWith(authStub.admin.id, { minimumFaceCount: 1 });
});
});
describe('getById', () => {
it('should throw a bad request when person is not found', async () => {
personMock.getById.mockResolvedValue(null);
await expect(sut.getById(authStub.admin, 'person-1')).rejects.toBeInstanceOf(BadRequestException);
});
it('should get a person by id', async () => {
personMock.getById.mockResolvedValue(personStub.withName);
await expect(sut.getById(authStub.admin, 'person-1')).resolves.toEqual(responseDto);
expect(personMock.getById).toHaveBeenCalledWith(authStub.admin.id, 'person-1');
});
});
describe('getThumbnail', () => {
it('should throw an error when personId is invalid', async () => {
personMock.getById.mockResolvedValue(null);
await expect(sut.getThumbnail(authStub.admin, 'person-1')).rejects.toBeInstanceOf(NotFoundException);
expect(storageMock.createReadStream).not.toHaveBeenCalled();
});
it('should throw an error when person has no thumbnail', async () => {
personMock.getById.mockResolvedValue(personStub.noThumbnail);
await expect(sut.getThumbnail(authStub.admin, 'person-1')).rejects.toBeInstanceOf(NotFoundException);
expect(storageMock.createReadStream).not.toHaveBeenCalled();
});
it('should serve the thumbnail', async () => {
personMock.getById.mockResolvedValue(personStub.noName);
await sut.getThumbnail(authStub.admin, 'person-1');
expect(storageMock.createReadStream).toHaveBeenCalledWith('/path/to/thumbnail', 'image/jpeg');
});
});
describe('getAssets', () => {
it("should return a person's assets", async () => {
personMock.getAssets.mockResolvedValue([assetEntityStub.image, assetEntityStub.video]);
await sut.getAssets(authStub.admin, 'person-1');
expect(personMock.getAssets).toHaveBeenCalledWith('admin_id', 'person-1');
});
});
describe('update', () => {
it('should throw an error when personId is invalid', async () => {
personMock.getById.mockResolvedValue(null);
await expect(sut.update(authStub.admin, 'person-1', { name: 'Person 1' })).rejects.toBeInstanceOf(
BadRequestException,
);
expect(personMock.update).not.toHaveBeenCalled();
});
it("should update a person's name", async () => {
personMock.getById.mockResolvedValue(personStub.noName);
personMock.update.mockResolvedValue(personStub.withName);
personMock.getAssets.mockResolvedValue([assetEntityStub.image]);
await expect(sut.update(authStub.admin, 'person-1', { name: 'Person 1' })).resolves.toEqual(responseDto);
expect(personMock.getById).toHaveBeenCalledWith('admin_id', 'person-1');
expect(personMock.update).toHaveBeenCalledWith({ id: 'person-1', name: 'Person 1' });
expect(jobMock.queue).toHaveBeenCalledWith({
name: JobName.SEARCH_INDEX_ASSET,
data: { ids: [assetEntityStub.image.id] },
});
});
});
describe('handlePersonCleanup', () => {
it('should delete people without faces', async () => {
personMock.getAllWithoutFaces.mockResolvedValue([personStub.noName]);
await sut.handlePersonCleanup();
expect(personMock.delete).toHaveBeenCalledWith(personStub.noName);
expect(jobMock.queue).toHaveBeenCalledWith({
name: JobName.DELETE_FILES,
data: { files: ['/path/to/thumbnail'] },
});
});
it('should log an error', async () => {
personMock.getAllWithoutFaces.mockResolvedValue([personStub.noName]);
personMock.delete.mockRejectedValue(new Error('database unavailable'));
await sut.handlePersonCleanup();
expect(jobMock.queue).not.toHaveBeenCalled();
});
});
});

View File

@@ -0,0 +1,82 @@
import { BadRequestException, Inject, Injectable, Logger, NotFoundException } from '@nestjs/common';
import { AssetResponseDto, mapAsset } from '../asset';
import { AuthUserDto } from '../auth';
import { IJobRepository, JobName } from '../job';
import { ImmichReadStream, IStorageRepository } from '../storage';
import { PersonUpdateDto } from './dto';
import { IPersonRepository } from './person.repository';
import { mapPerson, PersonResponseDto } from './response-dto';
@Injectable()
export class PersonService {
readonly logger = new Logger(PersonService.name);
constructor(
@Inject(IPersonRepository) private repository: IPersonRepository,
@Inject(IStorageRepository) private storageRepository: IStorageRepository,
@Inject(IJobRepository) private jobRepository: IJobRepository,
) {}
async getAll(authUser: AuthUserDto): Promise<PersonResponseDto[]> {
const people = await this.repository.getAll(authUser.id, { minimumFaceCount: 1 });
const named = people.filter((person) => !!person.name);
const unnamed = people.filter((person) => !person.name);
return (
[...named, ...unnamed]
// with thumbnails
.filter((person) => !!person.thumbnailPath)
.map((person) => mapPerson(person))
);
}
async getById(authUser: AuthUserDto, personId: string): Promise<PersonResponseDto> {
const person = await this.repository.getById(authUser.id, personId);
if (!person) {
throw new BadRequestException();
}
return mapPerson(person);
}
async getThumbnail(authUser: AuthUserDto, personId: string): Promise<ImmichReadStream> {
const person = await this.repository.getById(authUser.id, personId);
if (!person || !person.thumbnailPath) {
throw new NotFoundException();
}
return this.storageRepository.createReadStream(person.thumbnailPath, 'image/jpeg');
}
async getAssets(authUser: AuthUserDto, personId: string): Promise<AssetResponseDto[]> {
const assets = await this.repository.getAssets(authUser.id, personId);
return assets.map(mapAsset);
}
async update(authUser: AuthUserDto, personId: string, dto: PersonUpdateDto): Promise<PersonResponseDto> {
const exists = await this.repository.getById(authUser.id, personId);
if (!exists) {
throw new BadRequestException();
}
const person = await this.repository.update({ id: personId, name: dto.name });
const relatedAsset = await this.getAssets(authUser, personId);
const assetIds = relatedAsset.map((asset) => asset.id);
await this.jobRepository.queue({ name: JobName.SEARCH_INDEX_ASSET, data: { ids: assetIds } });
return mapPerson(person);
}
async handlePersonCleanup(): Promise<void> {
const people = await this.repository.getAllWithoutFaces();
for (const person of people) {
this.logger.debug(`Person ${person.name || person.id} no longer has any faces, deleting.`);
try {
await this.repository.delete(person);
await this.jobRepository.queue({ name: JobName.DELETE_FILES, data: { files: [person.thumbnailPath] } });
} catch (error: Error | any) {
this.logger.error(`Unable to delete person: ${error}`, error?.stack);
}
}
}
}

View File

@@ -0,0 +1 @@
export * from './person-response.dto';

View File

@@ -0,0 +1,19 @@
import { AssetFaceEntity, PersonEntity } from '@app/infra/entities';
export class PersonResponseDto {
id!: string;
name!: string;
thumbnailPath!: string;
}
export function mapPerson(person: PersonEntity): PersonResponseDto {
return {
id: person.id,
name: person.name,
thumbnailPath: person.thumbnailPath,
};
}
export function mapFace(face: AssetFaceEntity): PersonResponseDto {
return mapPerson(face.person);
}

View File

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

View File

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

View File

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

View File

@@ -4,9 +4,25 @@ export interface MachineLearningInput {
thumbnailPath: string;
}
export interface BoundingBox {
x1: number;
y1: number;
x2: number;
y2: number;
}
export interface DetectFaceResult {
imageWidth: number;
imageHeight: number;
boundingBox: BoundingBox;
score: number;
embedding: number[];
}
export interface IMachineLearningRepository {
classifyImage(input: MachineLearningInput): Promise<string[]>;
detectObjects(input: MachineLearningInput): Promise<string[]>;
encodeImage(input: MachineLearningInput): Promise<number[]>;
encodeText(input: string): Promise<number[]>;
detectFaces(input: MachineLearningInput): Promise<DetectFaceResult[]>;
}

View File

@@ -435,7 +435,7 @@ describe(UserService.name, () => {
{ deletedAt: makeDeletedAt(5) },
] as UserEntity[]);
await sut.handleQueueUserDelete();
await sut.handleUserDeleteCheck();
expect(userRepositoryMock.getDeletedUsers).toHaveBeenCalled();
expect(jobMock.queue).not.toHaveBeenCalled();
@@ -445,7 +445,7 @@ describe(UserService.name, () => {
const user = { deletedAt: makeDeletedAt(10) };
userRepositoryMock.getDeletedUsers.mockResolvedValue([user] as UserEntity[]);
await sut.handleQueueUserDelete();
await sut.handleUserDeleteCheck();
expect(userRepositoryMock.getDeletedUsers).toHaveBeenCalled();
expect(jobMock.queue).toHaveBeenCalledWith({ name: JobName.USER_DELETION, data: { user } });

View File

@@ -143,7 +143,7 @@ export class UserService {
return { admin, password, provided: !!providedPassword };
}
async handleQueueUserDelete() {
async handleUserDeleteCheck() {
const users = await this.userRepository.getDeletedUsers();
for (const user of users) {
if (this.isReadyForDeletion(user)) {

View File

@@ -0,0 +1,9 @@
import { IFaceRepository } from '../src';
export const newFaceRepositoryMock = (): jest.Mocked<IFaceRepository> => {
return {
getAll: jest.fn(),
getByIds: jest.fn(),
create: jest.fn(),
};
};

View File

@@ -3,6 +3,7 @@ import {
APIKeyEntity,
AssetEntity,
AssetType,
PersonEntity,
PartnerEntity,
SharedLinkEntity,
SharedLinkType,
@@ -10,6 +11,7 @@ import {
TranscodePreset,
UserEntity,
UserTokenEntity,
AssetFaceEntity,
} from '@app/infra/entities';
import {
AlbumResponseDto,
@@ -142,6 +144,7 @@ export const assetEntityStub = {
livePhotoVideoId: null,
tags: [],
sharedLinks: [],
faces: [],
}),
image: Object.freeze<AssetEntity>({
id: 'asset-id',
@@ -168,6 +171,7 @@ export const assetEntityStub = {
tags: [],
sharedLinks: [],
originalFileName: 'asset-id.ext',
faces: [],
}),
video: Object.freeze<AssetEntity>({
id: 'asset-id',
@@ -194,6 +198,7 @@ export const assetEntityStub = {
livePhotoVideoId: null,
tags: [],
sharedLinks: [],
faces: [],
}),
livePhotoMotionAsset: Object.freeze({
id: 'live-photo-motion-asset',
@@ -372,6 +377,7 @@ const assetResponse: AssetResponseDto = {
exifInfo: assetInfo,
livePhotoVideoId: null,
tags: [],
people: [],
};
const albumResponse: AlbumResponseDto = {
@@ -655,6 +661,7 @@ export const sharedLinkStub = {
},
tags: [],
sharedLinks: [],
faces: [],
},
],
},
@@ -729,6 +736,7 @@ export const searchStub = {
page: 1,
items: [],
facets: [],
distances: [],
}),
};
@@ -826,6 +834,39 @@ export const probeStub = {
}),
};
export const personStub = {
noName: Object.freeze<PersonEntity>({
id: 'person-1',
createdAt: new Date('2021-01-01'),
updatedAt: new Date('2021-01-01'),
ownerId: userEntityStub.admin.id,
owner: userEntityStub.admin,
name: '',
thumbnailPath: '/path/to/thumbnail',
faces: [],
}),
withName: Object.freeze<PersonEntity>({
id: 'person-1',
createdAt: new Date('2021-01-01'),
updatedAt: new Date('2021-01-01'),
ownerId: userEntityStub.admin.id,
owner: userEntityStub.admin,
name: 'Person 1',
thumbnailPath: '/path/to/thumbnail',
faces: [],
}),
noThumbnail: Object.freeze<PersonEntity>({
id: 'person-1',
createdAt: new Date('2021-01-01'),
updatedAt: new Date('2021-01-01'),
ownerId: userEntityStub.admin.id,
owner: userEntityStub.admin,
name: '',
thumbnailPath: '',
faces: [],
}),
};
export const partnerStub = {
adminToUser1: Object.freeze<PartnerEntity>({
createdAt: new Date('2023-02-23T05:06:29.716Z'),
@@ -844,3 +885,13 @@ export const partnerStub = {
sharedWith: userEntityStub.admin,
}),
};
export const faceStub = {
face1: Object.freeze<AssetFaceEntity>({
assetId: assetEntityStub.image.id,
asset: assetEntityStub.image,
personId: personStub.withName.id,
person: personStub.withName,
embedding: [1, 2, 3, 4],
}),
};

View File

@@ -3,11 +3,13 @@ export * from './api-key.repository.mock';
export * from './asset.repository.mock';
export * from './communication.repository.mock';
export * from './crypto.repository.mock';
export * from './face.repository.mock';
export * from './fixtures';
export * from './job.repository.mock';
export * from './machine-learning.repository.mock';
export * from './media.repository.mock';
export * from './partner.repository.mock';
export * from './person.repository.mock';
export * from './search.repository.mock';
export * from './shared-link.repository.mock';
export * from './smart-info.repository.mock';

View File

@@ -6,5 +6,6 @@ export const newMachineLearningRepositoryMock = (): jest.Mocked<IMachineLearning
detectObjects: jest.fn(),
encodeImage: jest.fn(),
encodeText: jest.fn(),
detectFaces: jest.fn(),
};
};

View File

@@ -5,6 +5,7 @@ export const newMediaRepositoryMock = (): jest.Mocked<IMediaRepository> => {
extractThumbnailFromExif: jest.fn(),
extractVideoThumbnail: jest.fn(),
resize: jest.fn(),
crop: jest.fn(),
probe: jest.fn(),
transcode: jest.fn(),
};

View File

@@ -0,0 +1,15 @@
import { IPersonRepository } from '../src';
export const newPersonRepositoryMock = (): jest.Mocked<IPersonRepository> => {
return {
getById: jest.fn(),
getAll: jest.fn(),
getAssets: jest.fn(),
getAllWithoutFaces: jest.fn(),
create: jest.fn(),
update: jest.fn(),
deleteAll: jest.fn(),
delete: jest.fn(),
};
};

View File

@@ -6,11 +6,15 @@ export const newSearchRepositoryMock = (): jest.Mocked<ISearchRepository> => {
checkMigrationStatus: jest.fn(),
importAssets: jest.fn(),
importAlbums: jest.fn(),
importFaces: jest.fn(),
deleteAlbums: jest.fn(),
deleteAssets: jest.fn(),
deleteFaces: jest.fn(),
deleteAllFaces: jest.fn(),
searchAssets: jest.fn(),
searchAlbums: jest.fn(),
vectorSearch: jest.fn(),
explore: jest.fn(),
searchFaces: jest.fn(),
};
};