diff --git a/web/packages/new/photos/services/ml/index.ts b/web/packages/new/photos/services/ml/index.ts index 6ad354a25a..846b48c62f 100644 --- a/web/packages/new/photos/services/ml/index.ts +++ b/web/packages/new/photos/services/ml/index.ts @@ -13,14 +13,14 @@ import { ComlinkWorker } from "@/base/worker/comlink-worker"; import { FileType } from "@/media/file-type"; import type { EnteFile } from "@/new/photos/types/file"; import { ensure } from "@/utils/ensure"; -import { throttled, wait } from "@/utils/promise"; +import { throttled } from "@/utils/promise"; import { proxy, transfer } from "comlink"; import { isInternalUser } from "../feature-flags"; import { getAllLocalFiles } from "../files"; import { getRemoteFlag, updateRemoteFlag } from "../remote-store"; import type { SearchPerson } from "../search/types"; import type { UploadItem } from "../upload/types"; -import { clusterFaces } from "./cluster-new"; +import { clusterFaces, type CGroup, type FaceCluster } from "./cluster-new"; import { regenerateFaceCrops } from "./crop"; import { clearMLDB, @@ -347,12 +347,21 @@ export const wipCluster = async () => { log.info("clustering"); _wip_isClustering = true; + _wip_searchPersons = undefined; triggerStatusUpdate(); - await wait(2000); - _wip_searchPersons = undefined; - const { clusters, cgroups } = await clusterFaces(await faceIndexes()); + const searchPersons = await convertToSearchPersons(clusters, cgroups); + + _wip_isClustering = false; + _wip_searchPersons = searchPersons; + triggerStatusUpdate(); +}; + +const convertToSearchPersons = async ( + clusters: FaceCluster[], + cgroups: CGroup[], +) => { const clusterByID = new Map(clusters.map((c) => [c.id, c])); const localFiles = await getAllLocalFiles(); @@ -360,43 +369,38 @@ export const wipCluster = async () => { const result: SearchPerson[] = []; for (const cgroup of cgroups) { - let avatarFaceID = cgroup.avatarFaceID; - // TODO-Cluster - // Temp - if (!avatarFaceID) { - // eslint-disable-next-line @typescript-eslint/no-non-null-assertion - avatarFaceID = cgroup.clusterIDs - .map((id) => clusterByID.get(id)) - .flatMap((cluster) => cluster?.faceIDs ?? [])[0]!; - } - cgroup.clusterIDs; - const avatarFaceFileID = fileIDFromFaceID(avatarFaceID); - const avatarFaceFile = localFileByID.get(avatarFaceFileID ?? 0); - if (!avatarFaceFileID || !avatarFaceFile) { - assertionFailed(`Face ID ${avatarFaceID} without local file`); + const displayFaceID = cgroup.displayFaceID; + if (!displayFaceID) { + // TODO-Cluster + assertionFailed(`cgroup ${cgroup.id} without displayFaceID`); continue; } - const files = cgroup.clusterIDs + + const displayFaceFileID = fileIDFromFaceID(displayFaceID); + if (!displayFaceFileID) continue; + + const displayFaceFile = localFileByID.get(displayFaceFileID); + if (!displayFaceFile) { + assertionFailed(`Face ID ${displayFaceFileID} without local file`); + continue; + } + + const fileIDs = cgroup.clusterIDs .map((id) => clusterByID.get(id)) .flatMap((cluster) => cluster?.faceIDs ?? []) .map((faceID) => fileIDFromFaceID(faceID)) .filter((fileID) => fileID !== undefined); + result.push({ id: cgroup.id, name: cgroup.name, - files, - displayFaceID: avatarFaceID, - displayFaceFile: avatarFaceFile, + files: [...new Set(fileIDs)], + displayFaceID, + displayFaceFile, }); } - const searchPersons = result.sort( - (a, b) => b.files.length - a.files.length, - ); - - _wip_isClustering = false; - _wip_searchPersons = searchPersons; - triggerStatusUpdate(); + return result.sort((a, b) => b.files.length - a.files.length); }; export type MLStatus = @@ -568,8 +572,7 @@ export const unidentifiedFaceIDs = async ( * Extract the fileID of the {@link EnteFile} to which the face belongs from its * faceID. */ -// TODO-Cluster: temporary export to supress linter -export const fileIDFromFaceID = (faceID: string) => { +const fileIDFromFaceID = (faceID: string) => { const fileID = parseInt(faceID.split("_")[0] ?? ""); if (isNaN(fileID)) { assertionFailed(`Ignoring attempt to parse invalid faceID ${faceID}`);