[desktop] ML touchups (#1777)

This commit is contained in:
Manav Rathi 2024-05-20 14:49:50 +05:30 committed by GitHub
commit 41b22abc66
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 98 additions and 64 deletions

View file

@ -14,7 +14,7 @@
"build:ci": "yarn build-renderer && tsc", "build:ci": "yarn build-renderer && tsc",
"build:quick": "yarn build-renderer && yarn build-main:quick", "build:quick": "yarn build-renderer && yarn build-main:quick",
"dev": "concurrently --kill-others --success first --names 'main,rndr' \"yarn dev-main\" \"yarn dev-renderer\"", "dev": "concurrently --kill-others --success first --names 'main,rndr' \"yarn dev-main\" \"yarn dev-renderer\"",
"dev-main": "tsc && electron app/main.js", "dev-main": "tsc && electron .",
"dev-renderer": "cd ../web && yarn install && yarn dev:photos", "dev-renderer": "cd ../web && yarn install && yarn dev:photos",
"postinstall": "electron-builder install-app-deps", "postinstall": "electron-builder install-app-deps",
"lint": "yarn prettier --check --log-level warn . && eslint --ext .ts src && yarn tsc", "lint": "yarn prettier --check --log-level warn . && eslint --ext .ts src && yarn tsc",

View file

@ -163,7 +163,7 @@ const checkForUpdatesAndNotify = async (mainWindow: BrowserWindow) => {
}; };
/** /**
* Return the version of the desktop app * Return the version of the desktop app.
* *
* The return value is of the form `v1.2.3`. * The return value is of the form `v1.2.3`.
*/ */

View file

@ -88,7 +88,8 @@ const fetchOrCreateImageBitmap = async (
const indexFaces_ = async (enteFile: EnteFile, imageBitmap: ImageBitmap) => { const indexFaces_ = async (enteFile: EnteFile, imageBitmap: ImageBitmap) => {
const fileID = enteFile.id; const fileID = enteFile.id;
const imageDimensions: Dimensions = imageBitmap; const { width, height } = imageBitmap;
const imageDimensions = { width, height };
const mlFile: MlFileData = { const mlFile: MlFileData = {
fileId: fileID, fileId: fileID,
mlVersion: defaultMLVersion, mlVersion: defaultMLVersion,
@ -126,8 +127,6 @@ const indexFaces_ = async (enteFile: EnteFile, imageBitmap: ImageBitmap) => {
const embeddings = await computeEmbeddings(alignedFacesData); const embeddings = await computeEmbeddings(alignedFacesData);
mlFile.faces.forEach((f, i) => (f.embedding = embeddings[i])); mlFile.faces.forEach((f, i) => (f.embedding = embeddings[i]));
// TODO-ML: Skip if somehow already relative. But why would it be?
// if (face.detection.box.x + face.detection.box.width < 2) continue;
mlFile.faces.forEach((face) => { mlFile.faces.forEach((face) => {
face.detection = relativeDetection(face.detection, imageDimensions); face.detection = relativeDetection(face.detection, imageDimensions);
}); });
@ -157,11 +156,6 @@ const detectFaces = async (
rect(imageBitmap), rect(imageBitmap),
); );
// TODO-ML: reenable faces filtering based on width ?? else remove me
// ?.filter((f) =>
// f.box.width > syncContext.config.faceDetection.minFaceSize
// );
const maxFaceDistancePercent = Math.sqrt(2) / 100; const maxFaceDistancePercent = Math.sqrt(2) / 100;
const maxFaceDistance = imageBitmap.width * maxFaceDistancePercent; const maxFaceDistance = imageBitmap.width * maxFaceDistancePercent;
return removeDuplicateDetections(faceDetections, maxFaceDistance); return removeDuplicateDetections(faceDetections, maxFaceDistance);
@ -320,8 +314,8 @@ const removeDuplicateDetections = (
const faceDetectionCenter = (detection: FaceDetection) => { const faceDetectionCenter = (detection: FaceDetection) => {
const center = new Point(0, 0); const center = new Point(0, 0);
// TODO-ML: first 4 landmarks is applicable to blazeface only this needs to // TODO-ML(LAURENS): first 4 landmarks is applicable to blazeface only this
// consider eyes, nose and mouth landmarks to take center // needs to consider eyes, nose and mouth landmarks to take center
detection.landmarks?.slice(0, 4).forEach((p) => { detection.landmarks?.slice(0, 4).forEach((p) => {
center.x += p.x; center.x += p.x;
center.y += p.y; center.y += p.y;
@ -354,11 +348,14 @@ const makeFaceID = (
const faceAlignment = (faceDetection: FaceDetection): FaceAlignment => const faceAlignment = (faceDetection: FaceDetection): FaceAlignment =>
faceAlignmentUsingSimilarityTransform( faceAlignmentUsingSimilarityTransform(
faceDetection, faceDetection,
normalizeLandmarks(arcFaceLandmarks, mobileFaceNetFaceSize), normalizeLandmarks(idealMobileFaceNetLandmarks, mobileFaceNetFaceSize),
); );
// TODO-ML: Rename? /**
const arcFaceLandmarks: [number, number][] = [ * The ideal location of the landmarks (eye etc) that the MobileFaceNet
* embedding model expects.
*/
const idealMobileFaceNetLandmarks: [number, number][] = [
[38.2946, 51.6963], [38.2946, 51.6963],
[73.5318, 51.5014], [73.5318, 51.5014],
[56.0252, 71.7366], [56.0252, 71.7366],
@ -681,21 +678,18 @@ const extractFaceCrop = (
imageBitmap: ImageBitmap, imageBitmap: ImageBitmap,
alignment: FaceAlignment, alignment: FaceAlignment,
): ImageBitmap => { ): ImageBitmap => {
// TODO-ML: Do we need to round twice? const alignmentBox = new Box({
const alignmentBox = roundBox( x: alignment.center.x - alignment.size / 2,
new Box({ y: alignment.center.y - alignment.size / 2,
x: alignment.center.x - alignment.size / 2, width: alignment.size,
y: alignment.center.y - alignment.size / 2, height: alignment.size,
width: alignment.size, });
height: alignment.size,
}),
);
const padding = 0.25; const padding = 0.25;
const scaleForPadding = 1 + padding * 2; const scaleForPadding = 1 + padding * 2;
const paddedBox = roundBox(enlargeBox(alignmentBox, scaleForPadding)); const paddedBox = roundBox(enlargeBox(alignmentBox, scaleForPadding));
// TODO-ML: The rotation doesn't seem to be used? it's set to 0. // TODO-ML(LAURENS): The rotation doesn't seem to be used? it's set to 0.
return cropWithRotation(imageBitmap, paddedBox, 0, 256); return cropWithRotation(imageBitmap, paddedBox, 0, 256);
}; };

View file

@ -12,15 +12,16 @@ export class DedicatedMLWorker {
public async syncLocalFile( public async syncLocalFile(
token: string, token: string,
userID: number, userID: number,
userAgent: string,
enteFile: EnteFile, enteFile: EnteFile,
localFile: globalThis.File, localFile: globalThis.File,
) { ) {
mlService.syncLocalFile(token, userID, enteFile, localFile); mlService.syncLocalFile(token, userID, userAgent, enteFile, localFile);
} }
public async sync(token: string, userID: number) { public async sync(token: string, userID: number, userAgent: string) {
await downloadManager.init(APPS.PHOTOS, { token }); await downloadManager.init(APPS.PHOTOS, { token });
return mlService.sync(token, userID); return mlService.sync(token, userID, userAgent);
} }
} }

View file

@ -24,7 +24,6 @@ export const syncPeopleIndex = async () => {
public async syncIndex(syncContext: MLSyncContext) { public async syncIndex(syncContext: MLSyncContext) {
await this.getMLLibraryData(syncContext); await this.getMLLibraryData(syncContext);
// TODO-ML(MR): Ensure this doesn't run until fixed.
await syncPeopleIndex(syncContext); await syncPeopleIndex(syncContext);
await this.persistMLLibraryData(syncContext); await this.persistMLLibraryData(syncContext);

View file

@ -8,8 +8,9 @@ import type { Face, FaceDetection, MlFileData } from "./types";
export const putFaceEmbedding = async ( export const putFaceEmbedding = async (
enteFile: EnteFile, enteFile: EnteFile,
mlFileData: MlFileData, mlFileData: MlFileData,
userAgent: string,
) => { ) => {
const serverMl = LocalFileMlDataToServerFileMl(mlFileData); const serverMl = LocalFileMlDataToServerFileMl(mlFileData, userAgent);
log.debug(() => ({ t: "Local ML file data", mlFileData })); log.debug(() => ({ t: "Local ML file data", mlFileData }));
log.debug(() => ({ log.debug(() => ({
t: "Uploaded ML file data", t: "Uploaded ML file data",
@ -57,34 +58,31 @@ class ServerFileMl {
class ServerFaceEmbeddings { class ServerFaceEmbeddings {
public faces: ServerFace[]; public faces: ServerFace[];
public version: number; public version: number;
/* TODO public client: string;
public client?: string;
public error?: boolean;
*/
public constructor(faces: ServerFace[], version: number) { public constructor(faces: ServerFace[], client: string, version: number) {
this.faces = faces; this.faces = faces;
this.client = client;
this.version = version; this.version = version;
} }
} }
class ServerFace { class ServerFace {
public faceID: string; public faceID: string;
// TODO-ML: singular? public embedding: number[];
public embeddings: number[];
public detection: ServerDetection; public detection: ServerDetection;
public score: number; public score: number;
public blur: number; public blur: number;
public constructor( public constructor(
faceID: string, faceID: string,
embeddings: number[], embedding: number[],
detection: ServerDetection, detection: ServerDetection,
score: number, score: number,
blur: number, blur: number,
) { ) {
this.faceID = faceID; this.faceID = faceID;
this.embeddings = embeddings; this.embedding = embedding;
this.detection = detection; this.detection = detection;
this.score = score; this.score = score;
this.blur = blur; this.blur = blur;
@ -122,6 +120,7 @@ class ServerFaceBox {
function LocalFileMlDataToServerFileMl( function LocalFileMlDataToServerFileMl(
localFileMlData: MlFileData, localFileMlData: MlFileData,
userAgent: string,
): ServerFileMl { ): ServerFileMl {
if (localFileMlData.errorCount > 0) { if (localFileMlData.errorCount > 0) {
return null; return null;
@ -140,7 +139,6 @@ function LocalFileMlDataToServerFileMl(
const landmarks = detection.landmarks; const landmarks = detection.landmarks;
const newBox = new ServerFaceBox(box.x, box.y, box.width, box.height); const newBox = new ServerFaceBox(box.x, box.y, box.width, box.height);
// TODO-ML: Add client UA and version
const newFaceObject = new ServerFace( const newFaceObject = new ServerFace(
faceID, faceID,
Array.from(embedding), Array.from(embedding),
@ -150,7 +148,7 @@ function LocalFileMlDataToServerFileMl(
); );
faces.push(newFaceObject); faces.push(newFaceObject);
} }
const faceEmbeddings = new ServerFaceEmbeddings(faces, 1); const faceEmbeddings = new ServerFaceEmbeddings(faces, userAgent, 1);
return new ServerFileMl( return new ServerFileMl(
localFileMlData.fileId, localFileMlData.fileId,
faceEmbeddings, faceEmbeddings,

View file

@ -1,9 +1,9 @@
import { Box, Point } from "services/face/geom"; import { Box, Point } from "services/face/geom";
import type { FaceDetection } from "services/face/types"; import type { FaceDetection } from "services/face/types";
// TODO-ML: Do we need two separate Matrix libraries? // TODO-ML(LAURENS): Do we need two separate Matrix libraries?
// //
// Keeping this in a separate file so that we can audit this. If these can be // Keeping this in a separate file so that we can audit this. If these can be
// expressed using ml-matrix, then we can move the code to f-index. // expressed using ml-matrix, then we can move this code to f-index.ts
import { import {
Matrix, Matrix,
applyToPoint, applyToPoint,

View file

@ -8,7 +8,7 @@ export interface FaceDetection {
} }
export interface FaceAlignment { export interface FaceAlignment {
// TODO-ML: remove affine matrix as rotation, size and center // TODO-ML(MR): remove affine matrix as rotation, size and center
// are simple to store and use, affine matrix adds complexity while getting crop // are simple to store and use, affine matrix adds complexity while getting crop
affineMatrix: number[][]; affineMatrix: number[][];
rotation: number; rotation: number;

View file

@ -11,11 +11,7 @@ import { EnteFile } from "types/file";
import { isInternalUserForML } from "utils/user"; import { isInternalUserForML } from "utils/user";
import { indexFaces } from "../face/f-index"; import { indexFaces } from "../face/f-index";
/** export const defaultMLVersion = 1;
* TODO-ML(MR): What and why.
* Also, needs to be 1 (in sync with mobile) when we move out of beta.
*/
export const defaultMLVersion = 3;
const batchSize = 200; const batchSize = 200;
@ -48,6 +44,7 @@ export async function updateMLSearchConfig(newConfig: MLSearchConfig) {
class MLSyncContext { class MLSyncContext {
public token: string; public token: string;
public userID: number; public userID: number;
public userAgent: string;
public localFilesMap: Map<number, EnteFile>; public localFilesMap: Map<number, EnteFile>;
public outOfSyncFiles: EnteFile[]; public outOfSyncFiles: EnteFile[];
@ -56,9 +53,10 @@ class MLSyncContext {
public syncQueue: PQueue; public syncQueue: PQueue;
constructor(token: string, userID: number) { constructor(token: string, userID: number, userAgent: string) {
this.token = token; this.token = token;
this.userID = userID; this.userID = userID;
this.userAgent = userAgent;
this.outOfSyncFiles = []; this.outOfSyncFiles = [];
this.nSyncedFiles = 0; this.nSyncedFiles = 0;
@ -81,12 +79,16 @@ class MachineLearningService {
private localSyncContext: Promise<MLSyncContext>; private localSyncContext: Promise<MLSyncContext>;
private syncContext: Promise<MLSyncContext>; private syncContext: Promise<MLSyncContext>;
public async sync(token: string, userID: number): Promise<boolean> { public async sync(
token: string,
userID: number,
userAgent: string,
): Promise<boolean> {
if (!token) { if (!token) {
throw Error("Token needed by ml service to sync file"); throw Error("Token needed by ml service to sync file");
} }
const syncContext = await this.getSyncContext(token, userID); const syncContext = await this.getSyncContext(token, userID, userAgent);
await this.syncLocalFiles(syncContext); await this.syncLocalFiles(syncContext);
@ -218,13 +220,17 @@ class MachineLearningService {
// await this.disposeMLModels(); // await this.disposeMLModels();
} }
private async getSyncContext(token: string, userID: number) { private async getSyncContext(
token: string,
userID: number,
userAgent: string,
) {
if (!this.syncContext) { if (!this.syncContext) {
log.info("Creating syncContext"); log.info("Creating syncContext");
// TODO-ML(MR): Keep as promise for now. // TODO-ML(MR): Keep as promise for now.
this.syncContext = new Promise((resolve) => { this.syncContext = new Promise((resolve) => {
resolve(new MLSyncContext(token, userID)); resolve(new MLSyncContext(token, userID, userAgent));
}); });
} else { } else {
log.info("reusing existing syncContext"); log.info("reusing existing syncContext");
@ -232,13 +238,17 @@ class MachineLearningService {
return this.syncContext; return this.syncContext;
} }
private async getLocalSyncContext(token: string, userID: number) { private async getLocalSyncContext(
token: string,
userID: number,
userAgent: string,
) {
// TODO-ML(MR): This is updating the file ML version. verify. // TODO-ML(MR): This is updating the file ML version. verify.
if (!this.localSyncContext) { if (!this.localSyncContext) {
log.info("Creating localSyncContext"); log.info("Creating localSyncContext");
// TODO-ML(MR): // TODO-ML(MR):
this.localSyncContext = new Promise((resolve) => { this.localSyncContext = new Promise((resolve) => {
resolve(new MLSyncContext(token, userID)); resolve(new MLSyncContext(token, userID, userAgent));
}); });
} else { } else {
log.info("reusing existing localSyncContext"); log.info("reusing existing localSyncContext");
@ -258,10 +268,15 @@ class MachineLearningService {
public async syncLocalFile( public async syncLocalFile(
token: string, token: string,
userID: number, userID: number,
userAgent: string,
enteFile: EnteFile, enteFile: EnteFile,
localFile?: globalThis.File, localFile?: globalThis.File,
) { ) {
const syncContext = await this.getLocalSyncContext(token, userID); const syncContext = await this.getLocalSyncContext(
token,
userID,
userAgent,
);
try { try {
await this.syncFileWithErrorHandler( await this.syncFileWithErrorHandler(
@ -285,7 +300,11 @@ class MachineLearningService {
localFile?: globalThis.File, localFile?: globalThis.File,
) { ) {
try { try {
const mlFileData = await this.syncFile(enteFile, localFile); const mlFileData = await this.syncFile(
enteFile,
localFile,
syncContext.userAgent,
);
syncContext.nSyncedFiles += 1; syncContext.nSyncedFiles += 1;
return mlFileData; return mlFileData;
} catch (e) { } catch (e) {
@ -317,14 +336,18 @@ class MachineLearningService {
} }
} }
private async syncFile(enteFile: EnteFile, localFile?: globalThis.File) { private async syncFile(
enteFile: EnteFile,
localFile: globalThis.File | undefined,
userAgent: string,
) {
const oldMlFile = await mlIDbStorage.getFile(enteFile.id); const oldMlFile = await mlIDbStorage.getFile(enteFile.id);
if (oldMlFile && oldMlFile.mlVersion) { if (oldMlFile && oldMlFile.mlVersion) {
return oldMlFile; return oldMlFile;
} }
const newMlFile = await indexFaces(enteFile, localFile); const newMlFile = await indexFaces(enteFile, localFile);
await putFaceEmbedding(enteFile, newMlFile); await putFaceEmbedding(enteFile, newMlFile, userAgent);
await mlIDbStorage.putFile(newMlFile); await mlIDbStorage.putFile(newMlFile);
return newMlFile; return newMlFile;
} }

View file

@ -1,6 +1,8 @@
import { FILE_TYPE } from "@/media/file-type"; import { FILE_TYPE } from "@/media/file-type";
import { ensureElectron } from "@/next/electron";
import log from "@/next/log"; import log from "@/next/log";
import { ComlinkWorker } from "@/next/worker/comlink-worker"; import { ComlinkWorker } from "@/next/worker/comlink-worker";
import { clientPackageNamePhotosDesktop } from "@ente/shared/apps/constants";
import { eventBus, Events } from "@ente/shared/events"; import { eventBus, Events } from "@ente/shared/events";
import { getToken, getUserID } from "@ente/shared/storage/localStorage/helpers"; import { getToken, getUserID } from "@ente/shared/storage/localStorage/helpers";
import debounce from "debounce"; import debounce from "debounce";
@ -227,8 +229,15 @@ class MLWorkManager {
this.stopSyncJob(); this.stopSyncJob();
const token = getToken(); const token = getToken();
const userID = getUserID(); const userID = getUserID();
const userAgent = await getUserAgent();
const mlWorker = await this.getLiveSyncWorker(); const mlWorker = await this.getLiveSyncWorker();
return mlWorker.syncLocalFile(token, userID, enteFile, localFile); return mlWorker.syncLocalFile(
token,
userID,
userAgent,
enteFile,
localFile,
);
}); });
} }
@ -266,9 +275,10 @@ class MLWorkManager {
const token = getToken(); const token = getToken();
const userID = getUserID(); const userID = getUserID();
const userAgent = await getUserAgent();
const jobWorkerProxy = await this.getSyncJobWorker(); const jobWorkerProxy = await this.getSyncJobWorker();
return await jobWorkerProxy.sync(token, userID); return await jobWorkerProxy.sync(token, userID, userAgent);
// this.terminateSyncJobWorker(); // this.terminateSyncJobWorker();
// TODO: redirect/refresh to gallery in case of session_expired, stop ml sync job // TODO: redirect/refresh to gallery in case of session_expired, stop ml sync job
} catch (e) { } catch (e) {
@ -320,3 +330,10 @@ export function logQueueStats(queue: PQueue, name: string) {
console.error(`queuestats: ${name}: Error, `, error), console.error(`queuestats: ${name}: Error, `, error),
); );
} }
const getUserAgent = async () => {
const electron = ensureElectron();
const name = clientPackageNamePhotosDesktop;
const version = await electron.appVersion();
return `${name}/${version}`;
};

View file

@ -14,6 +14,8 @@ export const CLIENT_PACKAGE_NAMES = new Map([
[APPS.ACCOUNTS, "io.ente.accounts.web"], [APPS.ACCOUNTS, "io.ente.accounts.web"],
]); ]);
export const clientPackageNamePhotosDesktop = "io.ente.photos.desktop";
export const APP_TITLES = new Map([ export const APP_TITLES = new Map([
[APPS.ALBUMS, "Ente Albums"], [APPS.ALBUMS, "Ente Albums"],
[APPS.PHOTOS, "Ente Photos"], [APPS.PHOTOS, "Ente Photos"],

View file

@ -28,8 +28,8 @@ class HTTPService {
const responseData = response.data; const responseData = response.data;
log.error( log.error(
`HTTP Service Error - ${JSON.stringify({ `HTTP Service Error - ${JSON.stringify({
url: config.url, url: config?.url,
method: config.method, method: config?.method,
xRequestId: response.headers["x-request-id"], xRequestId: response.headers["x-request-id"],
httpStatus: response.status, httpStatus: response.status,
errMessage: responseData.message, errMessage: responseData.message,