[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:quick": "yarn build-renderer && yarn build-main:quick",
"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",
"postinstall": "electron-builder install-app-deps",
"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`.
*/

View file

@ -88,7 +88,8 @@ const fetchOrCreateImageBitmap = async (
const indexFaces_ = async (enteFile: EnteFile, imageBitmap: ImageBitmap) => {
const fileID = enteFile.id;
const imageDimensions: Dimensions = imageBitmap;
const { width, height } = imageBitmap;
const imageDimensions = { width, height };
const mlFile: MlFileData = {
fileId: fileID,
mlVersion: defaultMLVersion,
@ -126,8 +127,6 @@ const indexFaces_ = async (enteFile: EnteFile, imageBitmap: ImageBitmap) => {
const embeddings = await computeEmbeddings(alignedFacesData);
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) => {
face.detection = relativeDetection(face.detection, imageDimensions);
});
@ -157,11 +156,6 @@ const detectFaces = async (
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 maxFaceDistance = imageBitmap.width * maxFaceDistancePercent;
return removeDuplicateDetections(faceDetections, maxFaceDistance);
@ -320,8 +314,8 @@ const removeDuplicateDetections = (
const faceDetectionCenter = (detection: FaceDetection) => {
const center = new Point(0, 0);
// TODO-ML: first 4 landmarks is applicable to blazeface only this needs to
// consider eyes, nose and mouth landmarks to take center
// TODO-ML(LAURENS): first 4 landmarks is applicable to blazeface only this
// needs to consider eyes, nose and mouth landmarks to take center
detection.landmarks?.slice(0, 4).forEach((p) => {
center.x += p.x;
center.y += p.y;
@ -354,11 +348,14 @@ const makeFaceID = (
const faceAlignment = (faceDetection: FaceDetection): FaceAlignment =>
faceAlignmentUsingSimilarityTransform(
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],
[73.5318, 51.5014],
[56.0252, 71.7366],
@ -681,21 +678,18 @@ const extractFaceCrop = (
imageBitmap: ImageBitmap,
alignment: FaceAlignment,
): ImageBitmap => {
// TODO-ML: Do we need to round twice?
const alignmentBox = roundBox(
new Box({
const alignmentBox = new Box({
x: alignment.center.x - alignment.size / 2,
y: alignment.center.y - alignment.size / 2,
width: alignment.size,
height: alignment.size,
}),
);
});
const padding = 0.25;
const scaleForPadding = 1 + padding * 2;
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);
};

View file

@ -12,15 +12,16 @@ export class DedicatedMLWorker {
public async syncLocalFile(
token: string,
userID: number,
userAgent: string,
enteFile: EnteFile,
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 });
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) {
await this.getMLLibraryData(syncContext);
// TODO-ML(MR): Ensure this doesn't run until fixed.
await syncPeopleIndex(syncContext);
await this.persistMLLibraryData(syncContext);

View file

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

View file

@ -1,9 +1,9 @@
import { Box, Point } from "services/face/geom";
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
// 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 {
Matrix,
applyToPoint,

View file

@ -8,7 +8,7 @@ export interface FaceDetection {
}
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
affineMatrix: number[][];
rotation: number;

View file

@ -11,11 +11,7 @@ import { EnteFile } from "types/file";
import { isInternalUserForML } from "utils/user";
import { indexFaces } from "../face/f-index";
/**
* 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;
export const defaultMLVersion = 1;
const batchSize = 200;
@ -48,6 +44,7 @@ export async function updateMLSearchConfig(newConfig: MLSearchConfig) {
class MLSyncContext {
public token: string;
public userID: number;
public userAgent: string;
public localFilesMap: Map<number, EnteFile>;
public outOfSyncFiles: EnteFile[];
@ -56,9 +53,10 @@ class MLSyncContext {
public syncQueue: PQueue;
constructor(token: string, userID: number) {
constructor(token: string, userID: number, userAgent: string) {
this.token = token;
this.userID = userID;
this.userAgent = userAgent;
this.outOfSyncFiles = [];
this.nSyncedFiles = 0;
@ -81,12 +79,16 @@ class MachineLearningService {
private localSyncContext: 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) {
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);
@ -218,13 +220,17 @@ class MachineLearningService {
// await this.disposeMLModels();
}
private async getSyncContext(token: string, userID: number) {
private async getSyncContext(
token: string,
userID: number,
userAgent: string,
) {
if (!this.syncContext) {
log.info("Creating syncContext");
// TODO-ML(MR): Keep as promise for now.
this.syncContext = new Promise((resolve) => {
resolve(new MLSyncContext(token, userID));
resolve(new MLSyncContext(token, userID, userAgent));
});
} else {
log.info("reusing existing syncContext");
@ -232,13 +238,17 @@ class MachineLearningService {
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.
if (!this.localSyncContext) {
log.info("Creating localSyncContext");
// TODO-ML(MR):
this.localSyncContext = new Promise((resolve) => {
resolve(new MLSyncContext(token, userID));
resolve(new MLSyncContext(token, userID, userAgent));
});
} else {
log.info("reusing existing localSyncContext");
@ -258,10 +268,15 @@ class MachineLearningService {
public async syncLocalFile(
token: string,
userID: number,
userAgent: string,
enteFile: EnteFile,
localFile?: globalThis.File,
) {
const syncContext = await this.getLocalSyncContext(token, userID);
const syncContext = await this.getLocalSyncContext(
token,
userID,
userAgent,
);
try {
await this.syncFileWithErrorHandler(
@ -285,7 +300,11 @@ class MachineLearningService {
localFile?: globalThis.File,
) {
try {
const mlFileData = await this.syncFile(enteFile, localFile);
const mlFileData = await this.syncFile(
enteFile,
localFile,
syncContext.userAgent,
);
syncContext.nSyncedFiles += 1;
return mlFileData;
} 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);
if (oldMlFile && oldMlFile.mlVersion) {
return oldMlFile;
}
const newMlFile = await indexFaces(enteFile, localFile);
await putFaceEmbedding(enteFile, newMlFile);
await putFaceEmbedding(enteFile, newMlFile, userAgent);
await mlIDbStorage.putFile(newMlFile);
return newMlFile;
}

View file

@ -1,6 +1,8 @@
import { FILE_TYPE } from "@/media/file-type";
import { ensureElectron } from "@/next/electron";
import log from "@/next/log";
import { ComlinkWorker } from "@/next/worker/comlink-worker";
import { clientPackageNamePhotosDesktop } from "@ente/shared/apps/constants";
import { eventBus, Events } from "@ente/shared/events";
import { getToken, getUserID } from "@ente/shared/storage/localStorage/helpers";
import debounce from "debounce";
@ -227,8 +229,15 @@ class MLWorkManager {
this.stopSyncJob();
const token = getToken();
const userID = getUserID();
const userAgent = await getUserAgent();
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 userID = getUserID();
const userAgent = await getUserAgent();
const jobWorkerProxy = await this.getSyncJobWorker();
return await jobWorkerProxy.sync(token, userID);
return await jobWorkerProxy.sync(token, userID, userAgent);
// this.terminateSyncJobWorker();
// TODO: redirect/refresh to gallery in case of session_expired, stop ml sync job
} catch (e) {
@ -320,3 +330,10 @@ export function logQueueStats(queue: PQueue, name: string) {
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"],
]);
export const clientPackageNamePhotosDesktop = "io.ente.photos.desktop";
export const APP_TITLES = new Map([
[APPS.ALBUMS, "Ente Albums"],
[APPS.PHOTOS, "Ente Photos"],

View file

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