Sztuczka z uruchamianiem rozpoznawania twarzy w Web Workerze w React

Kiedy chcemy wykonać proste zadania rozpoznawania twarzy na naszych zdjęciach, prawdopodobnie będziemy musieli, użyć aplikacji, która w tle wysyła zapytania do zewnętrznego API. Co w sytuacji gdy chemy zachować prywatność naszych danych i przetwarzać je lokalnie na komputerze? Sam byłem w takiej sytuacji kilka razy, ponieważ mam lekką obsesję na punkcie prywatności danych. Z punktu widzenia programisty JavaScript, jesli chodzi o budowanie takich aplikacji, zwykle myślimy o integrowaniu zewnętrznych usług lub prosimy kolegów znających Pythona o napisanie API (swoją drogą, to dobra okazja, żeby nauczyć się Pythona). Na szczęście istnieje dość stara, ale wciąż użyteczna biblioteka — face-api.js.

Czym jest face-api.js

Podczas wybierania biblioteki do użycia w projekcie jej wiek i zaangażowanie użytkowników w jej utrzymanie to zwykle jedne z głównych czynników. W tym przypadku jednak nie ma to większego znaczenia — face-api.js wydaje się być kompletną biblioteką, gotową do użycia w projekcie.

To rozwiązanie nie korzysta z LLM ani żadnych innych transformerów. face-api.js jest nakładką na tensorflow.js i współpracuje głównie z modelami: SSD Mobilenet V1, ResNet-34-like Face Recognition Model oraz MTCNN (experimental).

SSD (Single Shot MultiBox Detector) to model zaprojektowany do wykrywania wielu obiektów w jednym przebiegu, wytrenowany i dostosowany do zadań wykrywania twarzy. Model dobrze sprawdza się w zadaniach detekcji twarzy w czasie rzeczywistym. Jego rozmiar jest stosunkowo niewielki (około 5,4 MB), ponieważ został skwantyzowany.

ResNet-34-like Face Recognition Model został zaprojektowany do zadań rozpoznawania twarzy. Model ten potrafi generować deskryptory twarzy, które mogą zostać użyte do porównywania dwóch twarzy poprzez mierzenie ich podobieństwa z wykorzystaniem algorytmu Euclidean distance, którego implementacja jest również dostępna w face-api.js.

MTCNN (Multi-task Cascaded Convolutional Neural Networks) jest dołączony do face-api.js głównie w celach eksperymentalnych, ale może być bardzo przydatny. Model ten jest szczególnie wyspecjalizowany w zadaniach wykrywania twarzy. Jego główną cechą jest wykrywanie punktów charakterystycznych — na przykład oczu, ust i nosa. Wykrywanie odbywa się symultanicznie, Rozmiar tego modelu to około 2 MB.

Oto najbardziej, wedgług mnie, interesujące modele dostępne w face-api.js, ale jest ich więcej — na przykład lekka wersja do zadań rozpoznawania twarzy oraz model rozpoznawania ekspresji twarzy.

Oczywiście możliwe jest również użycie własnych modeli z face-api.js.

Z jakiego powodu używać `face-api.js` w Web Workerze?

Zanim zagłębimy się w implementację inferencjii w Web Workerze, warto zrozumieć, dlaczego opłaca się poświęcić wysiłek na zarządzanie komunikacją między głównym wątkiem a workerem.

Załóżmy, że mamy komponent, który pozwala użytkownikom przesłać przykładowe zdjęcie twarzy do przeglądarki, a następnie zestaw innych obrazów. Zadaniem komponentu jest anonimizacja każdej twarzy z podanego zestawu, która należy do osoby z przykładowego zdjęcia.

Najpierw musimy załadować modele:

typescript

import { useEffect } from "react";
import * as faceapi from "face-api.js";

const MODEL_PATH = `/models`;

async function loadModels(onError: (err: any) => void) {
  try {
    await Promise.all([
      faceapi.nets.ssdMobilenetv1.loadFromUri(MODEL_PATH),
      faceapi.nets.faceLandmark68Net.loadFromUri(MODEL_PATH),
      faceapi.nets.faceRecognitionNet.loadFromUri(MODEL_PATH),
    ]);
    console.log("models loaded");
  } catch (err) {
    onError("Failed to load face detection models");
    console.error(err);
  }
}

function useLoadModels() {
  useEffect(() => {
    loadModels((e) => console.log("Error Loading Models", e));
  }, []);
}

export default useLoadModels;

Modele możemy przechowywać w katalogu ./public. Po zamontowaniu komponentu wywoływana jest funkcja odpowiedzialna za ładowanie modeli.

Następnie zdefiniujemy naszą główną funkcję do inferencji. Na razie umieścimy całą logikę w hooku:

typescript

import { useState } from "react";
import * as faceapi from "face-api.js";

type ImageWithDescriptor = {
  id: number;
  descriptors: Float32Array<ArrayBufferLike>[];
  detections: faceapi.WithFaceDescriptor<
    faceapi.WithFaceLandmarks<
      {
        detection: faceapi.FaceDetection;
      },
      faceapi.FaceLandmarks68
    >
  >[];
  imgElement: HTMLCanvasElement;
};

async function createCanvasFromDataUrl(
  dataUrl: string,
): Promise<HTMLCanvasElement> {
  const res = await fetch(dataUrl);
  const blob = await res.blob();
  const bitmap = await createImageBitmap(blob);
  const canvas = document.createElement("canvas");
  canvas.width = bitmap.width;
  canvas.height = bitmap.height;
  const ctx = canvas.getContext("2d")!;
  ctx.drawImage(bitmap, 0, 0, canvas.width, canvas.height);
  return canvas;
}

function compareImages(
  exampleDescriptors: Float32Array[],
  targetImagesWithDescriptors: ImageWithDescriptor[],
) {
  const threshold = 0.5;

  const matchedImagesWithDescriptors: ImageWithDescriptor[] = [];
  targetImagesWithDescriptors.forEach(({ detections, ...rest }) => {
    const matchedDescriptors = detections.filter(({ descriptor }) => {
      return exampleDescriptors.some((exampleDescriptor) => {
        const distance = faceapi.euclideanDistance(
          exampleDescriptor,
          descriptor,
        );
        return distance < threshold;
      });
    });

    if (matchedDescriptors.length) {
      matchedImagesWithDescriptors.push({
        detections: matchedDescriptors,
        ...rest,
      });
    }
  });
  return matchedImagesWithDescriptors;
}

async function extractAllFaces(
  image: HTMLImageElement | HTMLCanvasElement | ImageBitmap,
) {
  const detections = await faceapi
    .detectAllFaces(image as any)
    .withFaceLandmarks()
    .withFaceDescriptors();

  return detections;
}

function useFace() {
  const [exampleImage, setExampleImage] = useState<string | null>(null);
  const [targetImages, setTargetImages] = useState<string[]>([]);
  const [isLoading, setIsLoading] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const [outputImages, setOutputImages] = useState<string[]>([]);

  const handleFace = async () => {
    if (!exampleImage) {
      return;
    }
    try {
      setIsLoading(true);
      setError(null);

      const exampleCanvas = await createCanvasFromDataUrl(exampleImage);

      const allFaces = await extractAllFaces(exampleCanvas);

      if (!allFaces.length) {
        setError("No faces detected in the example image");
        return;
      }

      const canvas = document.createElement("canvas");
      const ctx = canvas.getContext("2d")!;

      const updatedTargetImages: string[] = [];
      const targetImagesWithId = targetImages.map((targetImage, index) => ({
        id: index,
        src: targetImage,
      }));
      const targetImagesWithDescriptors = await Promise.all(
        targetImagesWithId.map(async ({ id, src }) => {
          const canvas = await createCanvasFromDataUrl(src);
          const detections = await extractAllFaces(canvas);

          return {
            id,
            detections,
            descriptors: detections.map((detection) => detection.descriptor),
            src,
            imgElement: canvas,
          };
        }),
      );

      const matchedImagesWithDescriptors = compareImages(
        allFaces.map((face) => face.descriptor),
        targetImagesWithDescriptors,
      );

      for (const targetImage of targetImagesWithDescriptors) {
        const detections =
          matchedImagesWithDescriptors.find(
            (matchedImageWithDescriptors) =>
              matchedImageWithDescriptors.id === targetImage.id,
          )?.detections || [];

        if (!detections?.length) {
          updatedTargetImages.push(targetImage.src);
          continue;
        }
        const targetImgElement = targetImage.imgElement;
        canvas.width = targetImgElement.width;
        canvas.height = targetImgElement.height;
        ctx.drawImage(targetImgElement, 0, 0, canvas.width, canvas.height);

        for (const detection of detections) {
          const { x, y, width, height } = detection?.detection?.box;
          ctx.filter = "blur(60px)";
          ctx.drawImage(
            targetImgElement,
            x,
            y,
            width,
            height,
            x,
            y,
            width,
            height,
          );
          ctx.filter = "none";
        }

        updatedTargetImages.push(canvas.toDataURL());
      }

      setOutputImages(updatedTargetImages);
    } catch (err) {
      setError("Error processing image");
      console.error(err);
    } finally {
      setIsLoading(false);
    }
  };

  return {
    handleFace,
    isLoading,
    error,
    outputImages,
    loadExampleImage: setExampleImage,
    loadTargetImages: setTargetImages,
    exampleImage,
    targetImages,
  };
}

export default useFace;

Tutaj dzieje się wiele rzeczy. Najważniejszą funkcją jest handleFace, która organizuje cały proces rozpoznawania, porównywania i blurrowania docelowych twarzy.

Najpierw musimy wyodrębnić wszystkie twarze z przykładowego zdjęcia. Twarze te posłużą jako przykłady do określenia, które twarze w docelowym zestawie należy rozmyć.

typescript

// REST OF THE CODE

async function extractAllFaces(image: HTMLCanvasElement) {
  const detections = await faceapi
    .detectAllFaces(image)
    .withFaceLandmarks()
    .withFaceDescriptors();

  return detections;
}

// REST OF THE CODE

const exampleCanvas = await createCanvasFromDataUrl(exampleImage);
const allFaces = await extractAllFaces(exampleCanvas);

// REST OF THE CODE

Następnie musimy wyodrębnić wszystkie twarze jako obiekty zwane deskryptorami z docelowego zestawu zdjęć, używając tego samego utilsa (extractAllFaces), co w poprzednim kroku.

typescript

// REST OF THE CODE
const targetImagesWithDescriptors = await Promise.all(
  targetImagesWithId.map(async ({ id, src }) => {
    const canvas = await createCanvasFromDataUrl(src);
    const detections = await extractAllFaces(canvas);

    return {
      id,
      detections,
      descriptors: detections.map((detection) => detection.descriptor),
      src,
      imgElement: canvas,
    };
  }),
); // REST OF THE CODE

Następnie porównujemy twarze, używając algorytmu Euclidean distance:

typescript

// REST OF THE CODE

function compareImages(
  exampleDescriptors: Float32Array[],
  targetImagesWithDescriptors: ImageWithDescriptor[],
) {
  const threshold = 0.5;

  const matchedImagesWithDescriptors: ImageWithDescriptor[] = [];
  targetImagesWithDescriptors.forEach(({ detections, ...rest }) => {
    const matchedDescriptors = detections.filter(({ descriptor }) => {
      return exampleDescriptors.some((exampleDescriptor) => {
        const distance = faceapi.euclideanDistance(
          exampleDescriptor,
          descriptor,
        );
        return distance < threshold;
      });
    });

    if (matchedDescriptors.length) {
      matchedImagesWithDescriptors.push({
        detections: matchedDescriptors,
        ...rest,
      });
    }
  });
  return matchedImagesWithDescriptors;
}

// REST OF THE CODE
const matchedImagesWithDescriptors = compareImages(
  allFaces.map((face) => face.descriptor),
  targetImagesWithDescriptors,
);
// REST OF THE CODE

face-api.js udostępnia również metodę do obliczania deskryptorów przy użyciu tego algorytmu.

Po pomyślnym porównaniu mamy wszystko, czego potrzebujemy, aby rozmyć odpowiednie twarze. Wykonujemy więc bardzo proste blurrowanie twarzy:

typescript

// REST OF THE CODE
for (const targetImage of targetImagesWithDescriptors) {
  const detections =
    matchedImagesWithDescriptors.find(
      (matchedImageWithDescriptors) =>
        matchedImageWithDescriptors.id === targetImage.id,
    )?.detections || [];

  if (!detections?.length) {
    updatedTargetImages.push(targetImage.src);
    continue;
  }
  const targetImgElement = targetImage.imgElement;
  canvas.width = targetImgElement.width;
  canvas.height = targetImgElement.height;
  ctx.drawImage(targetImgElement, 0, 0, canvas.width, canvas.height);

  for (const detection of detections) {
    const { x, y, width, height } = detection?.detection?.box;
    ctx.filter = "blur(60px)";
    ctx.drawImage(targetImgElement, x, y, width, height, x, y, width, height);
    ctx.filter = "none";
  }

  updatedTargetImages.push(canvas.toDataURL());
}
// REST OF THE CODE

Na koniec możemy wywołać hooki w komponencie React'owym:

JSX

import { useCallback } from "react";
import { Button } from "./ui/button";

import { Preview } from "./preview";
import { ImageUploader } from "./preview/image-uploader";
import useFace from "@/hooks/use-face";
import useLoadModels from "@/hooks/use-load-models";

function AnonymizerClient() {
  useLoadModels();
  const {
    outputImages,
    handleFace,
    loadExampleImage,
    loadTargetImages,
    isLoading,
    exampleImage,
    targetImages,
  } = useFace();

  const handleProcess = () => {
    handleFace();
  };

  const handleOnImageUpload = useCallback(
    (data: string[]) => {
      loadTargetImages((prevData) => [...prevData, ...data]);
    },
    [loadTargetImages],
  );
  return (
    <div className="flex w-full max-w-[800px] flex-col items-center gap-1.5">
      <Preview
        exampleImagePlaceholder={
          <ImageUploader
            type="single"
            onImageUpload={(img) => loadExampleImage(img[0])}
          />
        }
        targetImagesPlaceholder={
          <ImageUploader type="multiple" onImageUpload={handleOnImageUpload} />
        }
        images={outputImages?.length ? outputImages : targetImages}
        exampleImage={exampleImage}
      />
      <div className="flex space-x-4"></div>{" "}
      {isLoading && <span>Loading...</span>}
      <Button type="button" onClick={handleProcess} disabled={isLoading}>
        Process
      </Button>
    </div>
  );
}

export default AnonymizerClient;

Wygląda świetnie, więc…

…gdzie jest problem?

Problem polega na tym, że po kliknięciu przez użytkowników przycisku „Process”, interfejs użytkownika się zawiesza. Dzieje się tak, ponieważ środowisko uruchomieniowe JavaScriptu jest jednowątkowe. Domyślnie wszystko działa na jednym wątku procesora — tak zwanym głównym wątku — mimo że moglibyśmy mieć dostępnych nawet osiem lub więcej.

Mimo że modele są lekkie, połączenie zadań inferencji, renderowania interfejsu oraz rozmywania twarzy to zbyt duże obciążenie, by obsłużył je tylko główny wątek.

Rzowiązywanie problemu

Na szczęście istnieje rozwiązanie — Web Workery. Plan jest prosty: możemy przekazać kod odpowiedzialny za inferencję i blurrowanie do Web Workera. W ten sposób nasz interfejs pozostaje responsywny, ponieważ główny wątek będzie odciążony. W środowisku przeglądarki istnieje kilka typów workerów: Dedicated Worker, Shared Worker i Service Worker. Skupimy się na Dedicated Worker, który, według MDN, jest dostępny dla jednego dedykowanego skryptu, który go wywołał.

Aby z niego skorzystać, należy utworzyć plik zawierający kod workera:

typescript

// worker.js

onmessage(() => {
  console.log("Hello from worker");
  self.postMessage("Done");
});

Następnie, po stronie głównego wątku, należy utworzyć instancję klasy Worker, podając ścieżkę do pliku workera:

typescript

// main.js

const worker = new Worker("worker.js");

worker.onmessage = (event) => {
  console.log(event.data); // 'Hello from worker'
};

Komunikacja odbywa się za pomocą wiadomości.

Migracja logiki inferencji

Teraz, gdy problem został zidentyfikowany, logika inferencji musi zostać przeniesiona z głównego wątku do workera.

Do obsługi komunikacji użyjemy biblioteki, ponieważ natywne rozwiązania nie są zbyt wygodne. Narzędzie, którego użyjemy, nazywa się Comlink. Dzięki tej paczce wystarczy zdefiniować klasę, która będzie abstrakcją naszego workera; metody tej klasy są wywoływane, gdy odpowiednia wiadomość przychodzi z głównego wątku.

Oto przykład workera napisanego bez użycia jakichkolwiek bibliotek w celu zobrazowania :

typescript

// worker.ts

import type { MessageHandler } from "./types";

self.onmessage = (event: MessageEvent) => {
  const { type, payload } = event.data;

  switch (type) {
    case "printMessage":
      console.log(payload);
      break;

    case "getMessageLength":
      self.postMessage({
        type: "getMessageLengthResult",
        payload: payload.length,
      });
      break;
  }
};

// main.ts

import type { MessageHandler } from "./types";

const worker = new Worker(new URL("./worker.ts", import.meta.url), {
  type: "module",
});

const message = "Hello";

worker.onmessage = (event) => {
  const { type, payload } = event.data;

  if (type === "getMessageLengthResult") {
    console.log("Message length:", payload);
  }
};

worker.postMessage({ type: "printMessage", payload: message });
// logs -> 'Hello'
worker.postMessage({ type: "getMessageLength", payload: message });
// returns -> 5

Oto wersja wykorzystująca Comlink:

typescript

// worker.ts
import type { MessageHandler } trom './types'
import * as Comlink from "comlink";

class MessageWorker implements MessageHandler {
  printMessage(message: string) {
    console.log(message);
  }
  geMessageLength(message: string) {
    return message.length;
  }
}
const worker = new MessageWorker();
Comlink.expose(worker);

// main.ts
import type { MessageHandler } trom './types'
import * as Comlink from "comlink";

const worker = new Worker('worker.ts')
const workerApi = Comlink.wrap<MessageHandler>(worker)
const message = 'Hello'

workerApi.printMessage(message) // logs -> 'Hello'
workerApi.getMessageLength(message) // returns -> 5

);

Czytelniej, prawda?

Wszystko, co musimy zrobić, to przenieść naszą logikę inferencji do workera, prawda? Cóż... nie do końca.

Kolejny problem...

Jest jeden problem, który musimy rozwiązać, jeśli chcemy, aby nasza inferencja po stronie klienta była użyteczna.

Potrzebujemy canvas api, aby przetworzyć obraz, ale nie jest on dostępny w web workerze.

Problem z canvas jest również związany z face-api.js, ponieważ biblioteka korzysta z niego pod maską. Na szczęście istnieje zamiennik dla canvas, czyli OffscreenCanvas. Powinien również rozwiązać pozostałe problemy związane z blurrowaniem obrazu w web workerze.

Wreszcie możemy migrować

Najpierw stworzymy kod dla naszego workera. W ten sposób będziemy mogli zobaczyć, jak problem został rozwiązany.

Zdefiniujmy utilsy:

typescript

// serializers.ts

export function serializeDetection(det: any): any {
  if (!det) return det;

  const box = det.box ?? det._box;
  return {
    score: det.score ?? det._score,
    classScore: det.classScore ?? det._classScore,
    className: det.className ?? det._className,
    box: box ? serializeBox(box) : undefined,
    imageDims: det.imageDims ?? det._imageDims,
  };
}

export function serializeBox(box: any): any {
  return {
    x: box.x ?? box._x,
    y: box.y ?? box._y,
    width: box.width ?? box._width,
    height: box.height ?? box._height,
  };
}

export function serializeLandmarks(landmarks: any): any {
  if (!landmarks?.positions) return landmarks;
  return {
    positions: landmarks.positions.map((pt: any) => ({
      x: pt.x ?? pt._x,
      y: pt.y ?? pt._y,
    })),
  };
}
export function serializeFaceApiResult(result: any) {
  if (!result) return result;

  const out: Record<string, any> = {};

  if ("detection" in result) {
    out.detection = serializeDetection(result.detection);
  }
  if ("descriptor" in result) {
    out.descriptor = ArrayBuffer.isView(result.descriptor)
      ? Array.from(result.descriptor) // or keep Float32Array for zero-copy
      : result.descriptor;
  }
  if ("expression" in result) {
    out.expression = result.expression;
  }
  if ("landmarks" in result) {
    out.landmarks = serializeLandmarks(result.landmarks);
  }
  if ("alignedRect" in result) {
    out.alignedRect = serializeDetection(result.alignedRect);
  }

  return out;
}

Musimy zserializować właściwości obiektu, ponieważ obiekty face-api.js (detections, landmarks, descriptors) są złożonymi instancjami klas (obiektami z prototypami i metodami), za pomocą postMessage można przesyłać tylko serializowalne, proste typy danych JavaScript. Mówiąc prościej: jeśli nie zserializujemy tych obiektów, w głównym wątku otrzymalibyśmy brakujące lub niepoprawne właściwości. Dla przykładu box._y zamiast box.y.

Następnie konfigurujemy face-api.js w workerze i dodajemy również kilka utilsów:

typescript

// face-detection.worker.ts

import * as faceapi from "face-api.js";
import { serializeFaceApiResult } from "./worker/serializers";

const MODEL_PATH = `/models`;

faceapi.env.setEnv(faceapi.env.createNodejsEnv());

faceapi.env.monkeyPatch({
  Canvas: OffscreenCanvas,
  createCanvasElement: () => {
    return new OffscreenCanvas(480, 270);
  },
});

const createCanvas = async (transferObj: DataTransfer) => {
  try {
    const buf = transferObj as ArrayBuffer | undefined;
    if (!buf) {
      return new OffscreenCanvas(20, 20);
    }

    const blob = new Blob([buf]);
    const bitmap = await createImageBitmap(blob);
    const canvas = new OffscreenCanvas(bitmap.width, bitmap.height);
    const ctx = canvas.getContext("2d")!;
    ctx.drawImage(bitmap, 0, 0, bitmap.width, bitmap.height);
    return canvas;
  } catch (error) {
    console.error(
      "Error creating image from buffer, using empty canvas instead",
      error,
    );
    return new OffscreenCanvas(20, 20);
  }
};

Pierwszą rzeczą, którą należy zrobić, aby face-api.js działał w web workerze, jest jawne podanie alternatywy dla canvas. Mamy także funkcję do tworzenia canvas'a z danych obrazów, która będzie używana w wykrywaniu i rozpoznawaniu twarzy.

Teraz możemy zdefiniować naszą klasę, która dzięki Comlink będzie pełnić rolę „syntactic sugar” dla API workera:

typescript

class WorkerClass implements FaceDetectionWorker {
  async extractAllFaces(transferObj: DataTransfer) {
    const canvas = await createCanvas(transferObj);
    const detections = await faceapi
      .detectAllFaces(canvas as unknown as faceapi.TNetInput)
      .withFaceLandmarks()
      .withFaceDescriptors();
    return detections.map(serializeFaceApiResult);
  }

  async detectMatchingFaces(transferObj: {
    allExampleFaces: Float32Array[];
    allTargetImages: ArrayBuffer;
  }) {
    const allExampleFaces = transferObj.allExampleFaces;
    const detections = await this.extractAllFaces(transferObj.allTargetImages);
    const threshold = 0.5;

    const matchedDescriptors = detections.filter(({ descriptor }) => {
      return allExampleFaces.some((exampleDescriptor) => {
        const distance = faceapi.euclideanDistance(
          exampleDescriptor,
          descriptor,
        );
        return distance < threshold;
      });
    });
    return matchedDescriptors.map(serializeFaceApiResult);
  }

  async drawOutputImage(imageWithDescriptors: ImageWithDescriptors) {
    const decodedCanvas = await createCanvas(imageWithDescriptors.imgElement);

    const canvas = new OffscreenCanvas(
      decodedCanvas.width,
      decodedCanvas.height,
    );
    const ctxRes = canvas.getContext("2d", { willReadFrequently: true })!;
    const detections = imageWithDescriptors.detections;
    ctxRes.drawImage(
      decodedCanvas as unknown as CanvasImageSource,
      0,
      0,
      canvas.width,
      canvas.height,
    );

    for (const detection of detections) {
      const { x, y, width, height } = detection?.detection?.box;

      const padding = 0.2;
      const expandedX = Math.max(0, x - width * padding);
      const expandedY = Math.max(0, y - height * padding);
      const expandedWidth = Math.min(
        canvas.width - expandedX,
        width * (1 + 2 * padding),
      );
      const expandedHeight = Math.min(
        canvas.height - expandedY,
        height * (1 + 2 * padding),
      );

      for (let i = 0; i < 3; i++) {
        ctxRes.filter = "blur(50px)";
        ctxRes.drawImage(
          decodedCanvas as unknown as CanvasImageSource,
          expandedX,
          expandedY,
          expandedWidth,
          expandedHeight,
          expandedX,
          expandedY,
          expandedWidth,
          expandedHeight,
        );
      }
      ctxRes.filter = "none";
    }

    return canvas.convertToBlob();
  }
}

Podejście proponowane przez Comlink jest bardzo wygodne. Łatwo sobie wyobrazić, że nasze metody będą wywoływane w miejscach, gdzie worker odbiera wiadomości.

Przenieśliśmy również logikę obliczania porównań obrazów i twarzy, która wcześniej była wykonywana w funkcji handleFace oraz funkcji compareImages wywoływanej w useFaceHook. Teraz logika ta została zamknięta w metodzie detectMatchingFaces. Funkcja blurrowania również została przeniesiona i dostosowana.

Ostatnią rzeczą, którą musimy zrobić po stronie workera, jest załadowanie modeli i udostępnienie instancji naszej klasy workera:

typescript

// REST OF THE CODE

async function loadModels() {
  console.log("WorkerClass init...");
  await Promise.all([
    faceapi.nets.ssdMobilenetv1.loadFromUri(MODEL_PATH),
    faceapi.nets.faceLandmark68Net.loadFromUri(MODEL_PATH),
    faceapi.nets.faceRecognitionNet.loadFromUri(MODEL_PATH),
  ]);
  console.log("worker initialized and models loaded");
}

(async () => {
  await loadModels();
  const worker = new WorkerClass();
  Comlink.expose(worker);
})();

Przepisywanie logiki po stronie głównego wątku

Teraz musimy nieco dostosować logikę z hooka.

Pierwszą rzeczą, którą musimy dodać, jest inicjalizacja workera. Zainicjalizujemy workera w hooku useEffect i przypiszemy go do ref.

typescript

import { useEffect, useRef, useState } from "react";
import * as faceapi from "face-api.js";
import * as Comlink from "comlink";
import type { FaceDetectionWorker } from "../types";

function useFace() {
  const [exampleImage, setExampleImage] = useState<string | null>(null);
  const [targetImages, setTargetImages] = useState<string[]>([]);
  const [isLoading, setIsLoading] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const [outputImages, setOutputImages] = useState<string[]>([]);
  const workerRef = useRef<Worker | null>(null);

  useEffect(() => {
    const runWorker = async () => {
      const worker = new Worker(
        new URL("../face-detection.worker", import.meta.url),
        { type: "module" },
      );

      workerRef.current = worker;
    };

    runWorker().catch((error) => {
      console.error("Error initializing worker:", error);
    });

    return () => {
      workerRef.current?.terminate();
    };
  }, []);
}

// REST OF THE CODE...

Następnie, w naszym handlerze handleFace, musimy opakować worker, używając metody wrap z biblioteki Comlink.

typescript

function useFace() {
  // REST OF THE CODE

  const handleFace = async () => {
    if (!exampleImage || !workerRef.current) {
      return;
    }
    try {
      const api = Comlink.wrap<Comlink.Remote<FaceDetectionWorker>>(
        workerRef.current as any,
      );

      setIsLoading(true);
      setError(null);
    } catch (err) {
      setError("Error processing image");
      console.error(err);
    } finally {
      setIsLoading(false);
    }
  };

  // REST OF THE CODE
}

Dzięki temu otrzymujemy przyjazne i łatwo dostępne API, które pozwala nam płynnie komunikować się z workerem, tak jakby był klasą.

typescript

function useFace() {
  // REST OF THE CODE

  const handleFace = async () => {
    try {
      // REST OF THE CODE

      const exampleArrayBuffer = await (
        await fetch(exampleImage)
      ).arrayBuffer();

      const allExampleFaces = await api.extractAllFaces(exampleArrayBuffer);
      const targetImagesWithId = targetImages.map((targetImage, index) => ({
        id: index,
        src: targetImage,
      }));

      // REST OF THE CODE
    } catch (err) {
      setError("Error processing image");
      console.error(err);
    } finally {
      setIsLoading(false);
    }
  };

  // REST OF THE CODE
}

Po wykryciu wszystkich przykładowych twarzy na zdjęciu przykładowym możemy w końcu wykrywać i rozmywać twarze w docelowych obrazach.

typescript

// REST OF THE CODE

async function getImageWithDetections(
  allExampleFaces: Float32Array[],
  targetImageData: { id: number; src: string },
  detector: (transferObj: {
    allExampleFaces: Float32Array[];
    allTargetImages: ArrayBuffer;
  }) => Promise<
    faceapi.WithFaceDescriptor<
      faceapi.WithFaceLandmarks<
        { detection: faceapi.FaceDetection },
        faceapi.FaceLandmarks68
      >
    >[]
  >,
) {
  const { id, src } = targetImageData;
  const arrayBuffer = await (await fetch(src)).arrayBuffer();

  const arrayBufferForDetector = arrayBuffer.slice(0);

  const payload = {
    allExampleFaces,
    allTargetImages: arrayBufferForDetector,
  };

  const matchedDescriptors = await detector(payload);

  return {
    id,
    src,
    imgElement: arrayBuffer,
    detections: matchedDescriptors,
  };
}

function useFace() {
  // REST OF THE CODE

  const handleFace = async () => {
    try {
      // REST OF THE CODE
      for (const targetImage of targetImagesWithId) {
        const imageWithDescriptors = await getImageWithDetections(
          allExampleFaces.map((face) => face.descriptor),
          targetImage,
          api.detectMatchingFaces,
        );
        const output = await api.drawOutputImage(imageWithDescriptors);

        const url = URL.createObjectURL(output);
        setOutputImages((prevState) => [...prevState, url]);
      }
    } catch (err) {
      setError("Error processing image");
      console.error(err);
    } finally {
      setIsLoading(false);
    }
  };

  // REST OF THE CODE
}

Musimy utworzyć funkcję do uzyskiwania obrazu z wykrytymi twarzami, ponieważ nie chcemy zagnieżdżać zbyt dużo kodu w pętli — szybko stałoby się to chaotyczne. Funkcja przyjmuje wszystko, co potrzebne do uzyskania twarzy do rozmycia: przykładowe twarze, docelowy obraz, w którym szukamy twarzy pasujących do przykładów, oraz funkcję do wykrywania pasujących twarzy.

Cały moduł z hookiem useFace wygląda następująco:

typescript

import { useEffect, useRef, useState } from "react";
import * as faceapi from "face-api.js";
import * as Comlink from "comlink";
import type { FaceDetectionWorker } from "../types";

async function getImageWithDetections(
  allExampleFaces: Float32Array[],
  targetImageData: { id: number; src: string },
  detector: (transferObj: {
    allExampleFaces: Float32Array[];
    allTargetImages: ArrayBuffer;
  }) => Promise<
    faceapi.WithFaceDescriptor<
      faceapi.WithFaceLandmarks<
        { detection: faceapi.FaceDetection },
        faceapi.FaceLandmarks68
      >
    >[]
  >,
) {
  const { id, src } = targetImageData;
  const arrayBuffer = await (await fetch(src)).arrayBuffer();

  const arrayBufferForDetector = arrayBuffer.slice(0);

  const payload = {
    allExampleFaces,
    allTargetImages: arrayBufferForDetector,
  };

  const matchedDescriptors = await detector(payload);

  return {
    id,
    src,
    imgElement: arrayBuffer,
    detections: matchedDescriptors,
  };
}

function useFace() {
  const [exampleImage, setExampleImage] = useState<string | null>(null);
  const [targetImages, setTargetImages] = useState<string[]>([]);
  const [isLoading, setIsLoading] = useState(false);
  const [error, setError] = useState<string | null>(null);
  const [outputImages, setOutputImages] = useState<string[]>([]);
  const workerRef = useRef<Worker | null>(null);

  useEffect(() => {
    const runWorker = async () => {
      const worker = new Worker(
        new URL("../face-detection.worker", import.meta.url),
        { type: "module" },
      );

      workerRef.current = worker;
    };

    runWorker().catch((error) => {
      console.error("Error initializing worker:", error);
    });

    return () => {
      workerRef.current?.terminate();
    };
  }, []);

  const handleFace = async () => {
    if (!exampleImage || !workerRef.current) {
      return;
    }
    try {
      const api = Comlink.wrap<Comlink.Remote<FaceDetectionWorker>>(
        workerRef.current as any,
      );

      setIsLoading(true);
      setError(null);

      const exampleArrayBuffer = await (
        await fetch(exampleImage)
      ).arrayBuffer();

      const allExampleFaces = await api.extractAllFaces(exampleArrayBuffer);
      const targetImagesWithId = targetImages.map((targetImage, index) => ({
        id: index,
        src: targetImage,
      }));

      for (const targetImage of targetImagesWithId) {
        const imageWithDescriptors = await getImageWithDetections(
          allExampleFaces.map((face) => face.descriptor),
          targetImage,
          api.detectMatchingFaces,
        );
        const output = await api.drawOutputImage(imageWithDescriptors);

        const url = URL.createObjectURL(output);
        setOutputImages((prevState) => [...prevState, url]);
      }
    } catch (err) {
      setError("Error processing image");
      console.error(err);
    } finally {
      setIsLoading(false);
    }
  };

  return {
    handleFace,
    isLoading,
    error,
    outputImages,
    loadExampleImage: setExampleImage,
    loadTargetImages: setTargetImages,
    exampleImage,
    targetImages,
  };
}

export default useFace;

Podsumowanie

Po kilku trudnych krokach udało nam się rozwiązać problem, z którym wielu devów by się poddało i zwyczajnie skorzystało z zewnętrznych usług. W dzisiejszych czasach, gdy często jesteśmy połączeni z usługami w chmurze i zdalnymi serwerami, warto myśleć o użytkownikach i ich prywatności, również z naszej perspektywy może to być korzystne. Nie zawsze trzeba uruchamiać serwer albo funkcję serverless, przetwarzając dane użytkowników lokalnie na ich maszynie, nie trzeba się martwić o koszty wykorzystania zasobów. Ponadto nie ponosimy takiej samej odpowiedzialności za wrażliwe dane użytkowników przetwarzanych na ich maszynie, jak w przpadku przechowywania ich w naszej bazie danych na naszych serwerach.

Dzięki za przeczytanie! Stay tuned! 🫡

Sztuczka z uruchamianiem rozpoznawania twarzy w Web Workerze w React#

Czym jest face-api.js#

Z jakiego powodu używać face-api.js w Web Workerze?#

…gdzie jest problem?#

Rzowiązywanie problemu#

Migracja logiki inferencji#

Kolejny problem...#

Wreszcie możemy migrować#

Przepisywanie logiki po stronie głównego wątku#

Podsumowanie#

Sztuczka z uruchamianiem rozpoznawania twarzy w Web Workerze w React

Czym jest face-api.js

Z jakiego powodu używać `face-api.js` w Web Workerze?

…gdzie jest problem?

Rzowiązywanie problemu

Migracja logiki inferencji

Kolejny problem...

Wreszcie możemy migrować

Przepisywanie logiki po stronie głównego wątku

Podsumowanie