// demo下载 https://download.csdn.net/download/WWWLSW/88950321

// @canvas/image can decode jpeg, png, webp
const image = require('@canvas/image');
const fs = require('fs');
const tf = require('@tensorflow/tfjs-node');
const faceapi = require('@vladmandic/face-api');

const threshold = 0.4; // 可根据实际情况调整

async function compareFace() {
    const imageFile1 = 'wubai.png';
    const imageFile2 = 'renxianqi.png';
    const buffer1 = fs.readFileSync(imageFile1);
    const buffer2 = fs.readFileSync(imageFile2);
    const canvas1 = await image.imageFromBuffer(buffer1);
    const canvas2 = await image.imageFromBuffer(buffer2);
    const imageData1 = image.getImageData(canvas1);
    const imageData2 = image.getImageData(canvas2);

    const image1 = tf.tidy(() => {
        const data = tf.tensor(Array.from(imageData1?.data || []), [canvas1.height, canvas1.width, 4], 'int32');
        const channels = tf.split(data, 4, 2);
        const rgb = tf.stack([channels[0], channels[1], channels[2]], 2);
        const reshape = tf.reshape(rgb, [1, canvas1.height, canvas1.width, 3]);
        return reshape;
    });
    const image2 = tf.tidy(() => {
        const data = tf.tensor(Array.from(imageData2?.data || []), [canvas2.height, canvas2.width, 4], 'int32');
        const channels = tf.split(data, 4, 2);
        const rgb = tf.stack([channels[0], channels[1], channels[2]], 2);
        const reshape = tf.reshape(rgb, [1, canvas2.height, canvas2.width, 3]);
        return reshape;
    });

    // load models
    await faceapi.nets.ssdMobilenetv1.loadFromDisk("https://blog.csdn.net/WWWLSW/article/details/models");
    await faceapi.nets.ageGenderNet.loadFromDisk("https://blog.csdn.net/WWWLSW/article/details/models");
    await faceapi.nets.faceLandmark68Net.loadFromDisk("https://blog.csdn.net/WWWLSW/article/details/models");
    await faceapi.nets.faceRecognitionNet.loadFromDisk("https://blog.csdn.net/WWWLSW/article/details/models");
    console.log("begin:"+new Date())
    const face1 = await faceapi // run detection
        .detectSingleFace(image1)
        .withFaceLandmarks()
        .withAgeAndGender()
        .withFaceDescriptor();
    console.log("end1:"+new Date())
    const face2 = await faceapi // run detection
        .detectSingleFace(image2)
        .withFaceLandmarks()
        .withAgeAndGender()
        .withFaceDescriptor();
    console.log("end2:"+new Date())
    if (face1 && face2) {

        // Create an array of face descriptors
        const faceDescriptors = [face1.descriptor];

        // Create a FaceMatcher with the face descriptors
        const faceMatcher = new faceapi.FaceMatcher(faceDescriptors, threshold);

        // Compare the face descriptors of the second image
        const result = faceMatcher.findBestMatch(face2.descriptor);
        const person = result.label;
        if (person == "person 1") {
            console.log('这两张图片可能是同一个人。');
        } else {
            console.log('这两张图片可能是不同的人。');
        }
        // Output the result
        console.log(`Similarity: ${result.toString()}`);
        return result;
    } else {
        throw new Error('Unable to detect faces in one or both images.');
    }
}

compareFace();