github-mirrors/immich
1
0
Fork 0
mirror of https://github.com/immich-app/immich.git synced 2026-05-18 03:10:24 +03:00
Code Issues Packages Projects Releases Wiki Activity

refactor(server): job handlers (#2572)

Browse Source 
* refactor(server): job handlers

* chore: remove comment

* chore: add comments for
This commit is contained in:
Jason Rasmussen
2023-05-26 15:43:24 -04:00
committed by GitHub
parent d6756f3d81
commit 1c2d83e2c7
33 changed files with 807 additions and 1082 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/libs/domain/src/smart-info/smart-info.service.spec.ts
+25 -38
View File
@@ -30,6 +30,8 @@ describe(SmartInfoService.name, () => {
jobMock = newJobRepositoryMock();
machineMock = newMachineLearningRepositoryMock();
sut = new SmartInfoService(assetMock, jobMock, smartMock, machineMock);
assetMock.getByIds.mockResolvedValue([asset]);
});
it('should work', () => {
@@ -46,8 +48,8 @@ describe(SmartInfoService.name, () => {
await sut.handleQueueObjectTagging({ force: false });
expect(jobMock.queue.mock.calls).toEqual([
[{ name: JobName.CLASSIFY_IMAGE, data: { asset: assetEntityStub.image } }],
[{ name: JobName.DETECT_OBJECTS, data: { asset: assetEntityStub.image } }],
[{ name: JobName.CLASSIFY_IMAGE, data: { id: assetEntityStub.image.id } }],
[{ name: JobName.DETECT_OBJECTS, data: { id: assetEntityStub.image.id } }],
]);
expect(assetMock.getWithout).toHaveBeenCalledWith({ skip: 0, take: 1000 }, WithoutProperty.OBJECT_TAGS);
});
@@ -61,8 +63,8 @@ describe(SmartInfoService.name, () => {
await sut.handleQueueObjectTagging({ force: true });
expect(jobMock.queue.mock.calls).toEqual([
[{ name: JobName.CLASSIFY_IMAGE, data: { asset: assetEntityStub.image } }],
[{ name: JobName.DETECT_OBJECTS, data: { asset: assetEntityStub.image } }],
[{ name: JobName.CLASSIFY_IMAGE, data: { id: assetEntityStub.image.id } }],
[{ name: JobName.DETECT_OBJECTS, data: { id: assetEntityStub.image.id } }],
]);
expect(assetMock.getAll).toHaveBeenCalled();
});
@@ -70,7 +72,10 @@ describe(SmartInfoService.name, () => {
describe('handleTagImage', () => {
it('should skip assets without a resize path', async () => {
await sut.handleClassifyImage({ asset: { resizePath: '' } as AssetEntity });
const asset = { resizePath: '' } as AssetEntity;
assetMock.getByIds.mockResolvedValue([asset]);
await sut.handleClassifyImage({ id: asset.id });
expect(smartMock.upsert).not.toHaveBeenCalled();
expect(machineMock.classifyImage).not.toHaveBeenCalled();
@@ -79,7 +84,7 @@ describe(SmartInfoService.name, () => {
it('should save the returned tags', async () => {
machineMock.classifyImage.mockResolvedValue(['tag1', 'tag2', 'tag3']);
await sut.handleClassifyImage({ asset });
await sut.handleClassifyImage({ id: asset.id });
expect(machineMock.classifyImage).toHaveBeenCalledWith({ thumbnailPath: 'path/to/resize.ext' });
expect(smartMock.upsert).toHaveBeenCalledWith({
@@ -88,18 +93,10 @@ describe(SmartInfoService.name, () => {
});
});
it('should handle an error with the machine learning pipeline', async () => {
machineMock.classifyImage.mockRejectedValue(new Error('Unable to read thumbnail'));
await sut.handleClassifyImage({ asset });
expect(smartMock.upsert).not.toHaveBeenCalled();
});
it('should no update the smart info if no tags were returned', async () => {
machineMock.classifyImage.mockResolvedValue([]);
await sut.handleClassifyImage({ asset });
await sut.handleClassifyImage({ id: asset.id });
expect(machineMock.classifyImage).toHaveBeenCalled();
expect(smartMock.upsert).not.toHaveBeenCalled();
@@ -108,7 +105,10 @@ describe(SmartInfoService.name, () => {
describe('handleDetectObjects', () => {
it('should skip assets without a resize path', async () => {
await sut.handleDetectObjects({ asset: { resizePath: '' } as AssetEntity });
const asset = { resizePath: '' } as AssetEntity;
assetMock.getByIds.mockResolvedValue([asset]);
await sut.handleDetectObjects({ id: asset.id });
expect(smartMock.upsert).not.toHaveBeenCalled();
expect(machineMock.detectObjects).not.toHaveBeenCalled();
@@ -117,7 +117,7 @@ describe(SmartInfoService.name, () => {
it('should save the returned objects', async () => {
machineMock.detectObjects.mockResolvedValue(['obj1', 'obj2', 'obj3']);
await sut.handleDetectObjects({ asset });
await sut.handleDetectObjects({ id: asset.id });
expect(machineMock.detectObjects).toHaveBeenCalledWith({ thumbnailPath: 'path/to/resize.ext' });
expect(smartMock.upsert).toHaveBeenCalledWith({
@@ -126,18 +126,10 @@ describe(SmartInfoService.name, () => {
});
});
it('should handle an error with the machine learning pipeline', async () => {
machineMock.detectObjects.mockRejectedValue(new Error('Unable to read thumbnail'));
await sut.handleDetectObjects({ asset });
expect(smartMock.upsert).not.toHaveBeenCalled();
});
it('should no update the smart info if no objects were returned', async () => {
machineMock.detectObjects.mockResolvedValue([]);
await sut.handleDetectObjects({ asset });
await sut.handleDetectObjects({ id: asset.id });
expect(machineMock.detectObjects).toHaveBeenCalled();
expect(smartMock.upsert).not.toHaveBeenCalled();
@@ -153,7 +145,7 @@ describe(SmartInfoService.name, () => {
await sut.handleQueueEncodeClip({ force: false });
expect(jobMock.queue).toHaveBeenCalledWith({ name: JobName.ENCODE_CLIP, data: { asset: assetEntityStub.image } });
expect(jobMock.queue).toHaveBeenCalledWith({ name: JobName.ENCODE_CLIP, data: { id: assetEntityStub.image.id } });
expect(assetMock.getWithout).toHaveBeenCalledWith({ skip: 0, take: 1000 }, WithoutProperty.CLIP_ENCODING);
});
@@ -165,14 +157,17 @@ describe(SmartInfoService.name, () => {
await sut.handleQueueEncodeClip({ force: true });
expect(jobMock.queue).toHaveBeenCalledWith({ name: JobName.ENCODE_CLIP, data: { asset: assetEntityStub.image } });
expect(jobMock.queue).toHaveBeenCalledWith({ name: JobName.ENCODE_CLIP, data: { id: assetEntityStub.image.id } });
expect(assetMock.getAll).toHaveBeenCalled();
});
});
describe('handleEncodeClip', () => {
it('should skip assets without a resize path', async () => {
await sut.handleEncodeClip({ asset: { resizePath: '' } as AssetEntity });
const asset = { resizePath: '' } as AssetEntity;
assetMock.getByIds.mockResolvedValue([asset]);
await sut.handleEncodeClip({ id: asset.id });
expect(smartMock.upsert).not.toHaveBeenCalled();
expect(machineMock.encodeImage).not.toHaveBeenCalled();
@@ -181,7 +176,7 @@ describe(SmartInfoService.name, () => {
it('should save the returned objects', async () => {
machineMock.encodeImage.mockResolvedValue([0.01, 0.02, 0.03]);
await sut.handleEncodeClip({ asset });
await sut.handleEncodeClip({ id: asset.id });
expect(machineMock.encodeImage).toHaveBeenCalledWith({ thumbnailPath: 'path/to/resize.ext' });
expect(smartMock.upsert).toHaveBeenCalledWith({
@@ -189,13 +184,5 @@ describe(SmartInfoService.name, () => {
clipEmbedding: [0.01, 0.02, 0.03],
});
});
it('should handle an error with the machine learning pipeline', async () => {
machineMock.encodeImage.mockRejectedValue(new Error('Unable to read thumbnail'));
await sut.handleEncodeClip({ asset });
expect(smartMock.upsert).not.toHaveBeenCalled();
});
});
});
server/libs/domain/src/smart-info/smart-info.service.ts
+49 -58
View File
@@ -2,7 +2,7 @@ import { Inject, Injectable, Logger } from '@nestjs/common';
import { IAssetRepository, WithoutProperty } from '../asset';
import { MACHINE_LEARNING_ENABLED } from '../domain.constant';
import { usePagination } from '../domain.util';
import { IAssetJob, IBaseJob, IJobRepository, JobName, JOBS_ASSET_PAGINATION_SIZE } from '../job';
import { IBaseJob, IEntityJob, IJobRepository, JobName, JOBS_ASSET_PAGINATION_SIZE } from '../job';
import { IMachineLearningRepository } from './machine-learning.interface';
import { ISmartInfoRepository } from './smart-info.repository';
@@ -18,91 +18,82 @@ export class SmartInfoService {
) {}
async handleQueueObjectTagging({ force }: IBaseJob) {
try {
const assetPagination = usePagination(JOBS_ASSET_PAGINATION_SIZE, (pagination) => {
return force
? this.assetRepository.getAll(pagination)
: this.assetRepository.getWithout(pagination, WithoutProperty.OBJECT_TAGS);
});
const assetPagination = usePagination(JOBS_ASSET_PAGINATION_SIZE, (pagination) => {
return force
? this.assetRepository.getAll(pagination)
: this.assetRepository.getWithout(pagination, WithoutProperty.OBJECT_TAGS);
});
for await (const assets of assetPagination) {
for (const asset of assets) {
await this.jobRepository.queue({ name: JobName.CLASSIFY_IMAGE, data: { asset } });
await this.jobRepository.queue({ name: JobName.DETECT_OBJECTS, data: { asset } });
}
for await (const assets of assetPagination) {
for (const asset of assets) {
await this.jobRepository.queue({ name: JobName.CLASSIFY_IMAGE, data: { id: asset.id } });
await this.jobRepository.queue({ name: JobName.DETECT_OBJECTS, data: { id: asset.id } });
}
} catch (error: any) {
this.logger.error(`Unable to queue object tagging`, error?.stack);
}
return true;
}
async handleDetectObjects(data: IAssetJob) {
const { asset } = data;
async handleDetectObjects({ id }: IEntityJob) {
const [asset] = await this.assetRepository.getByIds([id]);
if (!MACHINE_LEARNING_ENABLED || !asset.resizePath) {
return;
return false;
}
try {
const objects = await this.machineLearning.detectObjects({ thumbnailPath: asset.resizePath });
if (objects.length > 0) {
await this.repository.upsert({ assetId: asset.id, objects });
await this.jobRepository.queue({ name: JobName.SEARCH_INDEX_ASSET, data: { ids: [asset.id] } });
}
} catch (error: any) {
this.logger.error(`Unable run object detection pipeline: ${asset.id}`, error?.stack);
const objects = await this.machineLearning.detectObjects({ thumbnailPath: asset.resizePath });
if (objects.length === 0) {
return false;
}
await this.repository.upsert({ assetId: asset.id, objects });
return true;
}
async handleClassifyImage(data: IAssetJob) {
const { asset } = data;
async handleClassifyImage({ id }: IEntityJob) {
const [asset] = await this.assetRepository.getByIds([id]);
if (!MACHINE_LEARNING_ENABLED || !asset.resizePath) {
return;
return false;
}
try {
const tags = await this.machineLearning.classifyImage({ thumbnailPath: asset.resizePath });
if (tags.length > 0) {
await this.repository.upsert({ assetId: asset.id, tags });
await this.jobRepository.queue({ name: JobName.SEARCH_INDEX_ASSET, data: { ids: [asset.id] } });
}
} catch (error: any) {
this.logger.error(`Unable to run image tagging pipeline: ${asset.id}`, error?.stack);
const tags = await this.machineLearning.classifyImage({ thumbnailPath: asset.resizePath });
if (tags.length === 0) {
return false;
}
await this.repository.upsert({ assetId: asset.id, tags });
return true;
}
async handleQueueEncodeClip({ force }: IBaseJob) {
try {
const assetPagination = usePagination(JOBS_ASSET_PAGINATION_SIZE, (pagination) => {
return force
? this.assetRepository.getAll(pagination)
: this.assetRepository.getWithout(pagination, WithoutProperty.CLIP_ENCODING);
});
const assetPagination = usePagination(JOBS_ASSET_PAGINATION_SIZE, (pagination) => {
return force
? this.assetRepository.getAll(pagination)
: this.assetRepository.getWithout(pagination, WithoutProperty.CLIP_ENCODING);
});
for await (const assets of assetPagination) {
for (const asset of assets) {
await this.jobRepository.queue({ name: JobName.ENCODE_CLIP, data: { asset } });
}
for await (const assets of assetPagination) {
for (const asset of assets) {
await this.jobRepository.queue({ name: JobName.ENCODE_CLIP, data: { id: asset.id } });
}
} catch (error: any) {
this.logger.error(`Unable to queue clip encoding`, error?.stack);
}
return true;
}
async handleEncodeClip(data: IAssetJob) {
const { asset } = data;
async handleEncodeClip({ id }: IEntityJob) {
const [asset] = await this.assetRepository.getByIds([id]);
if (!MACHINE_LEARNING_ENABLED || !asset.resizePath) {
return;
return false;
}
try {
const clipEmbedding = await this.machineLearning.encodeImage({ thumbnailPath: asset.resizePath });
await this.repository.upsert({ assetId: asset.id, clipEmbedding: clipEmbedding });
await this.jobRepository.queue({ name: JobName.SEARCH_INDEX_ASSET, data: { ids: [asset.id] } });
} catch (error: any) {
this.logger.error(`Unable run clip encoding pipeline: ${asset.id}`, error?.stack);
}
const clipEmbedding = await this.machineLearning.encodeImage({ thumbnailPath: asset.resizePath });
await this.repository.upsert({ assetId: asset.id, clipEmbedding: clipEmbedding });
return true;
}
}