aboutsummaryrefslogtreecommitdiff
path: root/src/client/views/search/FaceRecognitionHandler.tsx
blob: dc271fe73e21b45c1d18ded094eb2691787618c4 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
import * as faceapi from 'face-api.js';
import { FaceMatcher } from 'face-api.js';
import { computed } from 'mobx';
import { Doc, DocListCast } from '../../../fields/Doc';
import { DocData } from '../../../fields/DocSymbols';
import { List } from '../../../fields/List';
import { listSpec } from '../../../fields/Schema';
import { Cast, ImageCast, NumCast, StrCast } from '../../../fields/Types';
import { DocumentType } from '../../documents/DocumentTypes';
import { DocumentManager } from '../../util/DocumentManager';

/**
 * A class that handles face recognition.
 */
export class FaceRecognitionHandler {
    static _instance: FaceRecognitionHandler;
    private _loadedModels: boolean = false;
    private _processingDocs: Set<Doc> = new Set();
    private _pendingLoadDocs: Doc[] = [];

    public static FaceField = (target: Doc, doc: Doc) => `${Doc.LayoutFieldKey(target)}_${doc.face_label}`;
    public static FacesField = (target: Doc) => `${Doc.LayoutFieldKey(target)}_Faces`;

    constructor() {
        FaceRecognitionHandler._instance = this;
        this.loadModels().then(() => this._pendingLoadDocs.forEach(this.findMatches));
        DocumentManager.Instance.AddAnyViewRenderedCB(dv => FaceRecognitionHandler.Instance.findMatches(dv.Document));
    }

    @computed get examinedFaceDocs() {
        return DocListCast(Doc.UserDoc().examinedFaceDocs);
    }

    /**
     * Loads the face detection models.
     */
    loadModels = async () => {
        const MODEL_URL = `/models`;
        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();
    }

    /**
     * When a document is added, look for matching face documents.
     * @param doc The document being analyzed.
     */
    public findMatches = async (doc: Doc) => {
        if (!this._loadedModels || !Doc.ActiveDashboard) {
            this._pendingLoadDocs.push(doc);
            return;
        }

        if (doc.type === DocumentType.LOADING && !doc.loadingError) {
            setTimeout(() => this.findMatches(doc), 1000);
            return;
        }

        const imgUrl = ImageCast(doc[Doc.LayoutFieldKey(doc)]);
        // If the doc isn't an image or currently already been examined or is being processed, stop examining the document.
        if (!imgUrl || this.examinedFaceDocs.includes(doc) || this._processingDocs.has(doc)) {
            return;
        }

        // Mark the document as being processed.
        this._processingDocs.add(doc);

        // Get the image the document contains and analyze for faces.
        const [name, type] = imgUrl.url.href.split('.');
        const imageURL = `${name}_o.${type}`;

        const img = await this.loadImage(imageURL);

        const fullFaceDescriptions = await faceapi.detectAllFaces(img).withFaceLandmarks().withFaceDescriptors();

        doc[DocData][FaceRecognitionHandler.FacesField(doc)] = new List<List<number>>();

        // For each face detected, find a match.
        for (const fd of fullFaceDescriptions) {
            let match = this.findMatch(fd.descriptor);
            const converted_list = new List<number>(Array.from(fd.descriptor));

            if (match) {
                // If a matching Face Document has been found, add the document to the Face Document's associated docs and append the face
                // descriptor to the Face Document's descriptor list.
                Doc.AddDocToList(match, 'face_docList', doc);
                Cast(match.face_descriptors, listSpec('number'), null).push(converted_list as unknown as number); // items are lists of numbers, not numbers, but type system can't handle that
            } else {
                // If a matching Face Document has not been found, create a new Face Document.
                Doc.UserDoc().faceDocNum = NumCast(Doc.UserDoc().faceDocNum) + 1;

                const newFaceDocument = new Doc();
                newFaceDocument.title = `Face ${Doc.UserDoc().faceDocNum}`;
                newFaceDocument.face = ''; // just to make prettyprinting look better
                newFaceDocument.face_label = `Face${Doc.UserDoc().faceDocNum}`;
                newFaceDocument.face_docList = new List<Doc>([doc]);
                newFaceDocument.face_descriptors = new List<List<number>>([converted_list]);

                Doc.AddDocToList(Doc.ActiveDashboard[DocData], 'faceDocuments', newFaceDocument);
                match = newFaceDocument;
            }

            // Assign a field in the document of the matching Face Document.
            const faceDescripField = FaceRecognitionHandler.FaceField(doc, match);
            if (doc[DocData][faceDescripField]) {
                Cast(doc[DocData][faceDescripField], listSpec('number'), null).push(converted_list as unknown as number); // items are lists of numbers, not numbers, but type system can't handle that
            } else {
                doc[DocData][faceDescripField] = new List<List<number>>([converted_list]);
            }

            Cast(doc[DocData][FaceRecognitionHandler.FacesField(doc)], listSpec('number'), null).push(converted_list as unknown as number); // items are lists of numbers, not numbers, but type system can't handle that

            Doc.AddDocToList(Doc.UserDoc(), 'examinedFaceDocs', doc);
        }
        this._processingDocs.delete(doc);
    };

    /**
     * Finds a matching Face Document given a descriptor
     * @param cur_descriptor The current descriptor whose match is being searched for.
     * @returns The most similar Face Document.
     */
    private findMatch(cur_descriptor: Float32Array) {
        if (!Doc.ActiveDashboard || DocListCast(Doc.ActiveDashboard[DocData].faceDocuments).length < 1) {
            return null;
        }

        const faceDescriptors: faceapi.LabeledFaceDescriptors[] = DocListCast(Doc.ActiveDashboard[DocData].faceDocuments).map(faceDocument => {
            const float32Array = (faceDocument[DocData].face_descriptors as List<List<number>>).map(faceDescriptor => new Float32Array(Array.from(faceDescriptor)));
            return new faceapi.LabeledFaceDescriptors(StrCast(faceDocument[DocData].face_label), float32Array);
        });
        const faceMatcher = new FaceMatcher(faceDescriptors, 0.6);
        const match = faceMatcher.findBestMatch(cur_descriptor);
        if (match.label !== 'unknown') {
            for (const doc of DocListCast(Doc.ActiveDashboard[DocData].faceDocuments)) {
                if (doc[DocData].face_label === match.label) {
                    return doc;
                }
            }
        }
        return null;
    }

    /**
     * Loads an image
     */
    private loadImage = (src: string): Promise<HTMLImageElement> => {
        return new Promise((resolve, reject) => {
            const img = new Image();
            img.crossOrigin = 'anonymous';
            img.onload = () => resolve(img);
            img.onerror = err => reject(err);
            img.src = src;
        });
    };
}