From 320099df593c9555b70de58751153451401e41ff Mon Sep 17 00:00:00 2001
From: Shailesh Pandit <shailesh.pandit@gmail.com>
Date: Sat, 27 Nov 2021 15:31:03 +0530
Subject: [PATCH] Sync people from generated clusters Update personId in faces

---
 .../machineLearning/machineLearningService.ts | 74 +++++++++++++++----
 src/utils/machineLearning/index.ts            |  8 +-
 src/utils/machineLearning/types.ts            |  8 +-
 src/worker/machineLearning.worker.js          |  1 +
 4 files changed, 76 insertions(+), 15 deletions(-)

diff --git a/src/services/machineLearning/machineLearningService.ts b/src/services/machineLearning/machineLearningService.ts
index ce9bead87..bf7ae66ed 100644
--- a/src/services/machineLearning/machineLearningService.ts
+++ b/src/services/machineLearning/machineLearningService.ts
@@ -18,14 +18,16 @@ import {
     MLSyncConfig,
     MLSyncContext,
     MLSyncResult,
+    Person,
 } from 'utils/machineLearning/types';
 
 import * as jpeg from 'jpeg-js';
 import ClusteringService from './clusteringService';
 
 import { toTSNE } from 'utils/machineLearning/visualization';
-import { mlFilesStore } from 'utils/storage/localForage';
+import { mlFilesStore, mlPeopleStore } from 'utils/storage/localForage';
 import ArcfaceAlignmentService from './arcfaceAlignmentService';
+import { getAllFacesFromMap } from 'utils/machineLearning';
 
 class MachineLearningService {
     private faceDetectionService: FaceDetectionService;
@@ -140,7 +142,7 @@ class MachineLearningService {
             } catch (e) {
                 console.error(
                     'Error while syncing file: ',
-                    outOfSyncfile.id.toString(),
+                    outOfSyncfile.id,
                     e
                 );
             }
@@ -243,27 +245,38 @@ class MachineLearningService {
     }
 
     public async syncIndex(syncContext: MLSyncContext) {
-        await this.runFaceClustering(syncContext);
+        await this.syncPeopleIndex(syncContext);
     }
 
-    private async getAllFaces(syncContext: MLSyncContext) {
-        if (syncContext.allSyncedFaces) {
-            return syncContext.allSyncedFaces;
+    private async syncPeopleIndex(syncContext: MLSyncContext) {
+        const allFacesMap = await this.getAllSyncedFacesMap(syncContext);
+        const allFaces = getAllFacesFromMap(allFacesMap);
+
+        await this.runFaceClustering(syncContext, allFaces);
+        await this.syncPeopleFromClusters(syncContext, allFacesMap, allFaces);
+    }
+
+    private async getAllSyncedFacesMap(syncContext: MLSyncContext) {
+        if (syncContext.allSyncedFacesMap) {
+            return syncContext.allSyncedFacesMap;
         }
 
-        const allFaces: Array<Face> = [];
+        const allSyncedFacesMap = new Map<number, Array<Face>>();
         await mlFilesStore.iterate((mlFileData: MlFileData) => {
-            mlFileData.faces && allFaces.push(...mlFileData.faces);
+            mlFileData.faces &&
+                allSyncedFacesMap.set(mlFileData.fileId, mlFileData.faces);
         });
 
-        syncContext.allSyncedFaces = allFaces;
-        return allFaces;
+        syncContext.allSyncedFacesMap = allSyncedFacesMap;
+        return allSyncedFacesMap;
     }
 
-    public async runFaceClustering(syncContext: MLSyncContext) {
+    public async runFaceClustering(
+        syncContext: MLSyncContext,
+        allFaces: Array<Face>
+    ) {
         const clusteringConfig =
             syncContext.config.faceClustering.clusteringConfig;
-        const allFaces = await this.getAllFaces(syncContext);
 
         if (!allFaces || allFaces.length < clusteringConfig.minInputSize) {
             console.log(
@@ -294,10 +307,45 @@ class MachineLearningService {
         };
     }
 
+    private async syncPeopleFromClusters(
+        syncContext: MLSyncContext,
+        allFacesMap: Map<number, Array<Face>>,
+        allFaces: Array<Face>
+    ) {
+        const clusters = syncContext.faceClustersWithNoise?.clusters;
+        if (!clusters || clusters.length < 1) {
+            return;
+        }
+
+        await mlPeopleStore.clear();
+        for (const [index, cluster] of clusters.entries()) {
+            const faces = cluster.faces
+                .map((f) => allFaces[f])
+                .filter((f) => f);
+            const person: Person = {
+                personId: index,
+                files: faces.map((f) => f.fileId),
+            };
+
+            await mlPeopleStore.setItem(person.personId.toString(), person);
+
+            faces.forEach((face) => {
+                face.personId = person.personId;
+            });
+            // console.log("Creating person: ", person, faces);
+        }
+
+        await mlFilesStore.iterate((mlFileData: MlFileData, key) => {
+            mlFileData.faces = allFacesMap.get(mlFileData.fileId);
+            mlFilesStore.setItem(key, mlFileData);
+        });
+    }
+
     private async runTSNE(syncContext: MLSyncContext) {
         let faces = syncContext.syncedFaces;
         if (!faces || faces.length < 1) {
-            faces = await this.getAllFaces(syncContext);
+            const allFacesMap = await this.getAllSyncedFacesMap(syncContext);
+            faces = getAllFacesFromMap(allFacesMap);
         }
 
         const input = faces
diff --git a/src/utils/machineLearning/index.ts b/src/utils/machineLearning/index.ts
index c052c9c58..9ce38104e 100644
--- a/src/utils/machineLearning/index.ts
+++ b/src/utils/machineLearning/index.ts
@@ -1,7 +1,7 @@
 import * as tf from '@tensorflow/tfjs-core';
 import { DataType } from '@tensorflow/tfjs-core';
 import { Box, Point } from '../../../thirdparty/face-api/classes';
-import { MLSyncConfig } from './types';
+import { Face, MLSyncConfig } from './types';
 
 export function f32Average(descriptors: Float32Array[]) {
     if (descriptors.length < 1) {
@@ -130,6 +130,12 @@ export function computeRotation(point1: Point, point2: Point) {
     return normalizeRadians(radians);
 }
 
+export function getAllFacesFromMap(allFacesMap: Map<number, Array<Face>>) {
+    const allFaces = [...allFacesMap.values()].flat();
+
+    return allFaces;
+}
+
 export const DEFAULT_ML_SYNC_CONFIG: MLSyncConfig = {
     syncIntervalSec: 5, // 300
     batchSize: 5, // 200
diff --git a/src/utils/machineLearning/types.ts b/src/utils/machineLearning/types.ts
index cf8cbb841..2ee04657d 100644
--- a/src/utils/machineLearning/types.ts
+++ b/src/utils/machineLearning/types.ts
@@ -128,6 +128,12 @@ export interface Face extends FaceWithEmbedding {
     personId?: number;
 }
 
+export interface Person {
+    personId: number;
+    name?: string;
+    files: Array<number>;
+}
+
 export interface MlFileData {
     fileId: number;
     faces?: Face[];
@@ -184,7 +190,7 @@ export class MLSyncContext {
     outOfSyncFiles: File[];
     syncedFiles: File[];
     syncedFaces: Face[];
-    allSyncedFaces?: Face[];
+    allSyncedFacesMap?: Map<number, Array<Face>>;
     faceClusteringResults?: ClusteringResults;
     faceClustersWithNoise?: ClustersWithNoise;
     tsne?: any;
diff --git a/src/worker/machineLearning.worker.js b/src/worker/machineLearning.worker.js
index 1dff9d9b3..388d4466f 100644
--- a/src/worker/machineLearning.worker.js
+++ b/src/worker/machineLearning.worker.js
@@ -47,6 +47,7 @@ class MachineLearningWorker {
     }
 
     async sync(token) {
+        this.nextMLSyncTimeoutId = undefined;
         if (!runningInWorker()) {
             console.error(
                 'MachineLearning worker will only run in web worker env.'