From c82ef796f22c2c8ce232386589c7d4b5cc61a397 Mon Sep 17 00:00:00 2001 From: Manav Rathi Date: Mon, 8 Jul 2024 15:37:20 +0530 Subject: [PATCH] Extract --- web/packages/new/photos/services/ml/bitmap.ts | 28 ++++++++++ web/packages/new/photos/services/ml/face.ts | 55 ++++++++----------- web/packages/new/photos/services/ml/worker.ts | 8 ++- 3 files changed, 58 insertions(+), 33 deletions(-) diff --git a/web/packages/new/photos/services/ml/bitmap.ts b/web/packages/new/photos/services/ml/bitmap.ts index 3bcbb339ca..96741b2157 100644 --- a/web/packages/new/photos/services/ml/bitmap.ts +++ b/web/packages/new/photos/services/ml/bitmap.ts @@ -9,6 +9,34 @@ import DownloadManager from "../download"; import type { UploadItem } from "../upload/types"; import type { MLWorkerElectron } from "./worker-electron"; +export interface ImageBitmapAndData { + bitmap: ImageBitmap; +} + +/** + * Return an {@link ImageBitmap} and its {@link ImageData}. + * + * @param enteFile The {@link EnteFile} to index. + * + * @param uploadItem If we're called during the upload process, then this will + * be set to the {@link UploadItem} that was uploaded. This way, we can directly + * use the on-disk file instead of needing to download the original from remote. + * + * @param electron The {@link MLWorkerElectron} instance that allows us to call + * our Node.js layer for various functionality. + */ +export const imageBitmapAndData = async ( + enteFile: EnteFile, + uploadItem: UploadItem | undefined, + electron: MLWorkerElectron, +): Promise => { + const imageBitmap = uploadItem + ? await renderableUploadItemImageBitmap(enteFile, uploadItem, electron) + : await renderableImageBitmap(enteFile); + + return { bitmap: imageBitmap }; +}; + /** * Return a {@link ImageBitmap} that downloads the source image corresponding to * {@link enteFile} from remote. diff --git a/web/packages/new/photos/services/ml/face.ts b/web/packages/new/photos/services/ml/face.ts index 9f178f7b52..c90e7fd858 100644 --- a/web/packages/new/photos/services/ml/face.ts +++ b/web/packages/new/photos/services/ml/face.ts @@ -19,11 +19,7 @@ import { translate, type Matrix as TransformationMatrix, } from "transformation-matrix"; -import type { UploadItem } from "../upload/types"; -import { - renderableImageBitmap, - renderableUploadItemImageBitmap, -} from "./bitmap"; +import type { ImageBitmapAndData } from "./bitmap"; import { saveFaceCrops } from "./crop"; import { clamp, @@ -227,40 +223,37 @@ export interface Box { */ export const indexFaces = async ( enteFile: EnteFile, - uploadItem: UploadItem | undefined, + image: ImageBitmapAndData, electron: MLWorkerElectron, userAgent: string, ): Promise => { - const imageBitmap = uploadItem - ? await renderableUploadItemImageBitmap(enteFile, uploadItem, electron) - : await renderableImageBitmap(enteFile); + const imageBitmap = image.bitmap; + const { width, height } = imageBitmap; const fileID = enteFile.id; + const faceIndex = { + fileID, + width, + height, + faceEmbedding: { + version: faceIndexingVersion, + client: userAgent, + faces: await indexFacesInBitmap(fileID, imageBitmap, electron), + }, + }; + + // This step, saving face crops, is not part of the indexing pipeline; + // we just do it here since we have already have the ImageBitmap at + // hand. Ignore errors that happen during this since it does not impact + // the generated face index. try { - const faceIndex = { - fileID, - width, - height, - faceEmbedding: { - version: faceIndexingVersion, - client: userAgent, - faces: await indexFacesInBitmap(fileID, imageBitmap, electron), - }, - }; - // This step, saving face crops, is not part of the indexing pipeline; - // we just do it here since we have already have the ImageBitmap at - // hand. Ignore errors that happen during this since it does not impact - // the generated face index. - try { - await saveFaceCrops(imageBitmap, faceIndex); - } catch (e) { - log.error(`Failed to save face crops for file ${fileID}`, e); - } - return faceIndex; - } finally { - imageBitmap.close(); + await saveFaceCrops(imageBitmap, faceIndex); + } catch (e) { + log.error(`Failed to save face crops for file ${fileID}`, e); } + + return faceIndex; }; const indexFacesInBitmap = async ( diff --git a/web/packages/new/photos/services/ml/worker.ts b/web/packages/new/photos/services/ml/worker.ts index ccbac4d3b9..d854989c5f 100644 --- a/web/packages/new/photos/services/ml/worker.ts +++ b/web/packages/new/photos/services/ml/worker.ts @@ -11,6 +11,7 @@ import { expose } from "comlink"; import downloadManager from "../download"; import { getAllLocalFiles, getLocalTrashedFiles } from "../files"; import type { UploadItem } from "../upload/types"; +import { imageBitmapAndData } from "./bitmap"; import { indexableFileIDs, markIndexingFailed, @@ -328,7 +329,7 @@ const syncWithLocalFilesAndGetFilesToIndex = async ( * * @param userAgent The UA of the client that is doing the indexing (us). */ -export const index = async ( +const index = async ( enteFile: EnteFile, uploadItem: UploadItem | undefined, electron: MLWorkerElectron, @@ -337,13 +338,16 @@ export const index = async ( const f = fileLogID(enteFile); const startTime = Date.now(); + const image = await imageBitmapAndData(enteFile, uploadItem, electron); let faceIndex: FaceIndex; try { - faceIndex = await indexFaces(enteFile, uploadItem, electron, userAgent); + faceIndex = await indexFaces(enteFile, image, electron, userAgent); } catch (e) { log.error(`Failed to index faces in ${f}`, e); await markIndexingFailed(enteFile.id); throw e; + } finally { + image.bitmap.close(); } // [Note: Transient and permanent indexing failures]