feat(server): CLIP search integration (#1939)

This commit is contained in:
Alex
2023-03-18 08:44:42 -05:00
committed by GitHub
parent 0d436db3ea
commit f56eaae019
46 changed files with 673 additions and 773 deletions

View File

@@ -7,4 +7,6 @@ export interface MachineLearningInput {
export interface IMachineLearningRepository {
tagImage(input: MachineLearningInput): Promise<string[]>;
detectObjects(input: MachineLearningInput): Promise<string[]>;
encodeImage(input: MachineLearningInput): Promise<number[]>;
encodeText(input: string): Promise<number[]>;
}

View File

@@ -1,5 +1,6 @@
import { AssetEntity } from '@app/infra/db/entities';
import { newMachineLearningRepositoryMock, newSmartInfoRepositoryMock } from '../../test';
import { newJobRepositoryMock, newMachineLearningRepositoryMock, newSmartInfoRepositoryMock } from '../../test';
import { IJobRepository } from '../job';
import { IMachineLearningRepository } from './machine-learning.interface';
import { ISmartInfoRepository } from './smart-info.repository';
import { SmartInfoService } from './smart-info.service';
@@ -11,13 +12,15 @@ const asset = {
describe(SmartInfoService.name, () => {
let sut: SmartInfoService;
let jobMock: jest.Mocked<IJobRepository>;
let smartMock: jest.Mocked<ISmartInfoRepository>;
let machineMock: jest.Mocked<IMachineLearningRepository>;
beforeEach(async () => {
smartMock = newSmartInfoRepositoryMock();
jobMock = newJobRepositoryMock();
machineMock = newMachineLearningRepositoryMock();
sut = new SmartInfoService(smartMock, machineMock);
sut = new SmartInfoService(jobMock, smartMock, machineMock);
});
it('should work', () => {

View File

@@ -1,6 +1,6 @@
import { MACHINE_LEARNING_ENABLED } from '@app/common';
import { Inject, Injectable, Logger } from '@nestjs/common';
import { IAssetJob } from '../job';
import { IAssetJob, IJobRepository, JobName } from '../job';
import { IMachineLearningRepository } from './machine-learning.interface';
import { ISmartInfoRepository } from './smart-info.repository';
@@ -9,6 +9,7 @@ export class SmartInfoService {
private logger = new Logger(SmartInfoService.name);
constructor(
@Inject(IJobRepository) private jobRepository: IJobRepository,
@Inject(ISmartInfoRepository) private repository: ISmartInfoRepository,
@Inject(IMachineLearningRepository) private machineLearning: IMachineLearningRepository,
) {}
@@ -24,6 +25,7 @@ export class SmartInfoService {
const tags = await this.machineLearning.tagImage({ 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);
@@ -41,9 +43,26 @@ export class SmartInfoService {
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);
}
}
async handleEncodeClip(data: IAssetJob) {
const { asset } = data;
if (!MACHINE_LEARNING_ENABLED || !asset.resizePath) {
return;
}
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);
}
}
}