feat: integrated face recognition code with image documents

2 files changed, 65 insertions, 27 deletions

diff --git a/src/client/util/DocumentManager.ts b/src/client/util/DocumentManager.ts
index 8ad6ddf47..7d4684f41 100644
--- a/src/client/util/DocumentManager.ts
+++ b/src/client/util/DocumentManager.ts
@@ -6,11 +6,12 @@ import { Id } from '../../fields/FieldSymbols';
 import { listSpec } from '../../fields/Schema';
 import { Cast, DocCast, NumCast, StrCast } from '../../fields/Types';
 import { AudioField } from '../../fields/URLField';
-import { CollectionViewType } from '../documents/DocumentTypes';
+import { CollectionViewType, DocumentType } from '../documents/DocumentTypes';
 import { DocumentView, DocumentViewInternal } from '../views/nodes/DocumentView';
 import { FocusViewOptions } from '../views/nodes/FocusViewOptions';
 import { OpenWhere } from '../views/nodes/OpenWhere';
 import { PresBox } from '../views/nodes/trails';
+import { FaceRecognitionHandler } from '../views/search/FaceRecognitionHandler';
 
 type childIterator = { viewSpec: Opt<Doc>; childDocView: Opt<DocumentView>; focused: boolean; contextPath: Doc[] };
 export class DocumentManager {
@@ -96,6 +97,9 @@ export class DocumentManager {
 
     @action
     public AddView = (view: DocumentView) => {
+        if (view.Document.type === DocumentType.IMG && view.Document.embedContainer) {
+            FaceRecognitionHandler.Instance.findMatches(view.Document);
+        }
         this.AddDocumentView(view);
         this.callAddViewFuncs(view);
     };
diff --git a/src/client/views/search/FaceRecognitionHandler.tsx b/src/client/views/search/FaceRecognitionHandler.tsx
index 86619b2d1..fcd38c42f 100644
--- a/src/client/views/search/FaceRecognitionHandler.tsx
+++ b/src/client/views/search/FaceRecognitionHandler.tsx
@@ -4,51 +4,84 @@ import { Doc, DocListCast, NumListCast } from '../../../fields/Doc';
 import { DocData } from '../../../fields/DocSymbols';
 import { List } from '../../../fields/List';
 import { ObjectField } from '../../../fields/ObjectField';
-import { StrCast } from '../../../fields/Types';
+import { ImageCast, StrCast } from '../../../fields/Types';
+import { DocUtils } from '../../documents/DocUtils';
+import { Deserializable } from '../../util/SerializationHelper';
+import { DocumentView } from '../nodes/DocumentView';
 
 export class FaceRecognitionHandler {
     static _instance: FaceRecognitionHandler;
+    private loadedModels: boolean = false;
+    examinedDocs: Set<Doc> = new Set();
+    processingDocs: Set<Doc> = new Set();
 
     constructor() {
         FaceRecognitionHandler._instance = this;
         this.loadModels();
-        if (!Doc.ActiveDashboard![DocData].faceDocuments) {
-            Doc.ActiveDashboard![DocData].faceDocuments = new List<Doc>();
-        }
     }
 
     async loadModels() {
         const MODEL_URL = `/models`;
-        await faceapi.loadTinyFaceDetectorModel(MODEL_URL);
-        await faceapi.loadFaceLandmarkTinyModel(MODEL_URL);
+        await faceapi.loadFaceDetectionModel(MODEL_URL);
+        await faceapi.loadFaceLandmarkModel(MODEL_URL);
         await faceapi.loadFaceRecognitionModel(MODEL_URL);
+        this.loadedModels = true;
     }
 
     public static get Instance() {
         return FaceRecognitionHandler._instance ?? new FaceRecognitionHandler();
     }
 
-    public async findMatches(doc: Doc, imageURL: string) {
-        const img = await this.loadImage(imageURL);
-
-        const fullFaceDescriptions = await faceapi.detectAllFaces(img, new TinyFaceDetectorOptions()).withFaceLandmarks(true).withFaceDescriptors();
-
-        fullFaceDescriptions.forEach(fd => {
-            const match = this.findMatch(fd.descriptor);
-            if (match) {
-                match[DocData].associatedDocs = new List<Doc>([...DocListCast(match[DocData].associatedDocs), doc]);
-                match[DocData].faceDescriptors = new List<List<number>>([...(match[DocData].faceDescriptors as List<List<number>>), Array.from(fd.descriptor) as List<number>]);
-            } else {
-                const newFaceDocument = new Doc();
-                const converted_array = Array.from(fd.descriptor);
-                newFaceDocument[DocData].faceDescriptors = new List<List<number>>();
-                (newFaceDocument[DocData].faceDescriptors as List<List<number>>).push(converted_array as List<number>);
-                newFaceDocument[DocData].label = `Person ${DocListCast(Doc.ActiveDashboard![DocData].faceDocuments).length + 1}`;
-                newFaceDocument[DocData].associatedDocs = new List<Doc>([doc]);
-
-                Doc.ActiveDashboard![DocData].faceDocuments = new List<Doc>([...DocListCast(Doc.ActiveDashboard![DocData].faceDocuments), newFaceDocument]);
+    public async findMatches(doc: Doc) {
+        if (this.loadedModels) {
+            if (!Doc.ActiveDashboard![DocData].faceDocuments) {
+                Doc.ActiveDashboard![DocData].faceDocuments = new List<Doc>();
             }
-        });
+
+            if (this.examinedDocs.has(doc) || this.processingDocs.has(doc)) {
+                return;
+            }
+
+            this.processingDocs.add(doc);
+
+            try {
+                const [name, type] = ImageCast(doc[Doc.LayoutFieldKey(doc)]).url.href.split('.');
+                const imageURL = `${name}_o.${type}`;
+
+                const img = await this.loadImage(imageURL);
+
+                const fullFaceDescriptions = await faceapi.detectAllFaces(img).withFaceLandmarks().withFaceDescriptors();
+
+                for (const fd of fullFaceDescriptions) {
+                    const match = this.findMatch(fd.descriptor);
+                    if (match) {
+                        const converted_array = Array.from(fd.descriptor);
+                        const converted_list = new List<number>(converted_array);
+                        match[DocData].associatedDocs = new List<Doc>([...DocListCast(match[DocData].associatedDocs), doc]);
+                        match[DocData].faceDescriptors = new List<List<number>>([...(match[DocData].faceDescriptors as List<List<number>>), converted_list]);
+                    } else {
+                        const newFaceDocument = new Doc();
+                        const converted_array = Array.from(fd.descriptor);
+                        const converted_list = new List<number>(converted_array);
+                        newFaceDocument[DocData].faceDescriptors = new List<List<number>>();
+                        (newFaceDocument[DocData].faceDescriptors as List<List<number>>).push(converted_list);
+                        newFaceDocument[DocData].label = `Person ${DocListCast(Doc.ActiveDashboard![DocData].faceDocuments).length + 1}`;
+                        newFaceDocument[DocData].associatedDocs = new List<Doc>([doc]);
+
+                        Doc.ActiveDashboard![DocData].faceDocuments = new List<Doc>([...DocListCast(Doc.ActiveDashboard![DocData].faceDocuments), newFaceDocument]);
+                    }
+                }
+
+                this.examinedDocs.add(doc);
+                console.log(this.examinedDocs);
+
+                DocListCast(Doc.ActiveDashboard![DocData].faceDocuments).forEach(doc => console.log(DocListCast(doc[DocData].associatedDocs)));
+            } catch (error) {
+                console.error('Error processing document:', error);
+            } finally {
+                this.processingDocs.delete(doc);
+            }
+        }
     }
 
     private findMatch(cur_descriptor: Float32Array) {
@@ -68,6 +101,7 @@ export class FaceRecognitionHandler {
         } else {
             for (const doc of DocListCast(Doc.ActiveDashboard![DocData].faceDocuments)) {
                 if (doc[DocData].label === match.label) {
+                    console.log(match.label);
                     return doc;
                 }
             }