remote mapping

This commit is contained in:
Manav Rathi
2024-05-30 15:44:25 +05:30
parent 2b3b84de0f
commit 9dbec2729c
4 changed files with 32 additions and 154 deletions

View File

@@ -14,7 +14,7 @@ import {
translate,
} from "transformation-matrix";
import type { EnteFile } from "types/file";
import { getRenderableImage, logIdentifier } from "utils/file";
import { fileLogID, getRenderableImage } from "utils/file";
import { saveFaceCrop } from "./crop";
import {
clamp,
@@ -22,8 +22,7 @@ import {
pixelRGBBilinear,
warpAffineFloat32List,
} from "./image";
import type { Box, Dimensions, Face, Point } from "./types";
import type { MlFileData } from "./types-old";
import type { Box, Dimensions, Face, FaceIndex, Point } from "./types";
/**
* Index faces in the given file.
@@ -60,19 +59,20 @@ export const indexFaces = async (
const imageBitmap = await renderableImageBlob(enteFile, file).then(
createImageBitmap,
);
let mlFile: MlFileData;
let index: FaceIndex;
try {
mlFile = await indexFaces_(enteFile, imageBitmap, userAgent);
index = await indexFaces_(enteFile, imageBitmap, userAgent);
} finally {
imageBitmap.close();
}
log.debug(() => {
const nf = mlFile.faceEmbedding.faces?.length ?? 0;
const nf = index.faceEmbedding.faces.length;
const ms = Date.now() - startTime;
return `Indexed ${nf} faces in file ${logIdentifier(enteFile)} (${ms} ms)`;
return `Indexed ${nf} faces in ${fileLogID(enteFile)} (${ms} ms)`;
});
return mlFile;
return index;
};
/**
@@ -145,9 +145,9 @@ const indexFacesInBitmap = async (
const alignment = computeFaceAlignment(detection);
alignments.push(alignment);
// This step is not really part of the indexing pipeline, we just do
// it here since we have already computed the face alignment. Ignore
// errors that happen during this though.
// This step is not part of the indexing pipeline, we just do it here
// since we have already computed the face alignment. Ignore errors that
// happen during this since it does not impact the generated face index.
try {
await saveFaceCrop(imageBitmap, faceID, alignment);
} catch (e) {

View File

@@ -1,10 +1,14 @@
import log from "@/next/log";
import { putFaceEmbedding } from "services/face/remote";
import type { EnteFile } from "types/file";
import { logIdentifier } from "utils/file";
import { closeFaceDBConnectionsIfNeeded, markIndexingFailed } from "./db";
import { fileLogID } from "utils/file";
import {
closeFaceDBConnectionsIfNeeded,
markIndexingFailed,
saveFaceIndex,
} from "./db";
import { indexFaces } from "./f-index";
import type { MlFileData } from "./types-old";
import { putFaceIndex } from "./remote";
import type { FaceIndex } from "./types";
/**
* Index faces in a file, save the persist the results locally, and put them on
@@ -29,22 +33,21 @@ export class FaceIndexerWorker {
* downloaded and decrypted from remote.
*/
async index(enteFile: EnteFile, file: File | undefined, userAgent: string) {
const f = logIdentifier(enteFile);
let faceIndex: MlFileData;
let faceIndex: FaceIndex;
try {
faceIndex = await indexFaces(enteFile, file, userAgent);
log.debug(() => ({ f, faceIndex }));
} catch (e) {
// Mark indexing as having failed only if the indexing itself
// failed, not if there were subsequent failures (like when trying
// to put the result to remote or save it to the local face DB).
log.error(`Failed to index faces in file ${f}`, e);
log.error(`Failed to index faces in ${fileLogID(enteFile)}`, e);
markIndexingFailed(enteFile.id);
throw e;
}
await putFaceEmbedding(enteFile, faceIndex);
await putFaceIndex(enteFile, faceIndex);
await saveFaceIndex(faceIndex);
return faceIndex;
}

View File

@@ -2,23 +2,20 @@ import log from "@/next/log";
import ComlinkCryptoWorker from "@ente/shared/crypto";
import { putEmbedding } from "services/embeddingService";
import type { EnteFile } from "types/file";
import type { Point } from "./types";
import type { MlFileData } from "./types-old";
import type { FaceIndex } from "./types";
export const putFaceEmbedding = async (
export const putFaceIndex = async (
enteFile: EnteFile,
mlFileData: MlFileData,
faceIndex: FaceIndex,
) => {
const serverMl = LocalFileMlDataToServerFileMl(mlFileData);
log.debug(() => ({ t: "Local ML file data", mlFileData }));
log.debug(() => ({
t: "Uploaded ML file data",
d: JSON.stringify(serverMl),
t: "Uploading faceEmbedding",
d: JSON.stringify(faceIndex),
}));
const comlinkCryptoWorker = await ComlinkCryptoWorker.getInstance();
const { file: encryptedEmbeddingData } =
await comlinkCryptoWorker.encryptMetadata(serverMl, enteFile.key);
await comlinkCryptoWorker.encryptMetadata(faceIndex, enteFile.key);
await putEmbedding({
fileID: enteFile.id,
encryptedEmbedding: encryptedEmbeddingData.encryptedData,
@@ -26,125 +23,3 @@ export const putFaceEmbedding = async (
model: "file-ml-clip-face",
});
};
export interface FileML extends ServerFileMl {
updatedAt: number;
}
class ServerFileMl {
public fileID: number;
public height?: number;
public width?: number;
public faceEmbedding: ServerFaceEmbedding;
public constructor(
fileID: number,
faceEmbedding: ServerFaceEmbedding,
height?: number,
width?: number,
) {
this.fileID = fileID;
this.height = height;
this.width = width;
this.faceEmbedding = faceEmbedding;
}
}
class ServerFaceEmbedding {
public faces: ServerFace[];
public version: number;
public client: string;
public constructor(faces: ServerFace[], client: string, version: number) {
this.faces = faces;
this.client = client;
this.version = version;
}
}
class ServerFace {
public faceID: string;
public embedding: number[];
public detection: ServerDetection;
public score: number;
public blur: number;
public constructor(
faceID: string,
embedding: number[],
detection: ServerDetection,
score: number,
blur: number,
) {
this.faceID = faceID;
this.embedding = embedding;
this.detection = detection;
this.score = score;
this.blur = blur;
}
}
class ServerDetection {
public box: ServerFaceBox;
public landmarks: Point[];
public constructor(box: ServerFaceBox, landmarks: Point[]) {
this.box = box;
this.landmarks = landmarks;
}
}
class ServerFaceBox {
public x: number;
public y: number;
public width: number;
public height: number;
public constructor(x: number, y: number, width: number, height: number) {
this.x = x;
this.y = y;
this.width = width;
this.height = height;
}
}
function LocalFileMlDataToServerFileMl(
localFileMlData: MlFileData,
): ServerFileMl {
if (localFileMlData.errorCount > 0) {
return null;
}
const faces: ServerFace[] = [];
for (let i = 0; i < localFileMlData.faceEmbedding.faces.length; i++) {
const face = localFileMlData.faceEmbedding.faces[i];
const faceID = face.faceID;
const embedding = face.embedding;
const score = face.score;
const blur = face.blur;
const detection = face.detection;
const box = detection.box;
const landmarks = detection.landmarks;
faces.push(
new ServerFace(
faceID,
Array.from(embedding),
new ServerDetection(box, landmarks),
score,
blur,
),
);
}
const faceEmbedding = new ServerFaceEmbedding(
faces,
localFileMlData.faceEmbedding.client,
localFileMlData.faceEmbedding.version,
);
return new ServerFileMl(
localFileMlData.fileID,
faceEmbedding,
localFileMlData.height,
localFileMlData.width,
);
}

View File

@@ -86,8 +86,8 @@ const moduleState = new ModuleState();
* given {@link enteFile}. The returned string contains the file name (for ease
* of debugging) and the file ID (for exactness).
*/
export const logIdentifier = (enteFile: EnteFile) =>
`${enteFile.metadata.title ?? "-"} (${enteFile.id})`;
export const fileLogID = (enteFile: EnteFile) =>
`file ${enteFile.metadata.title ?? "-"} (${enteFile.id})`;
export async function getUpdatedEXIFFileForDownload(
fileReader: FileReader,