Определение лица и маски
В предыдущей статье я рассказывал о том, можно ли использовать машинное обучение (в частности определение лица и маски) в браузере, подходах к детекции и оптимизации всех процессов.
Сегодня я хочу рассказать о технических подробностях реализации.


Использованные технологии




Основной язык для разработки в браузере это TypeScript. Клиентское приложение написано на React.js.

В приложении используется несколько нейронных сетей для детекции разных событий: детекция лица, детекция маски. Каждая модель/сеть запускается в отдельном потоке (Web Worker). Нейронные сети запускаются с использованием TensorFlow.js и в качестве backend-а используется Web Assembly или WebGL, что позволяет выполнять код со скоростью близкой к нативной. Выбор того или иного backend-а зависит от размера модели (мелкие модели быстрее работают на WebAssembly), но надо всегда проводить тестирование и выбирать, то что быстрее для конкретной модели.

Получение и отображение видео стрима с использованием WebRTC. Для работы с изображениями используется библиотека OpenCV.js.

Реализован был следующий подход:



Основной поток занимается только оркестрацией всего, он не загружает тяжелую библиотеку OpenCV для работы с изображениями и не использует TensorFlow.js. Все что он делает, получает изображения из видео потока и отправляет на обработку веб воркерам.

Пока воркер не сообщил основному потоку, что он освободился, новое изображение не посылается в него, тем самым не создается очередь, как только воркер говорит, что он освободился, текущее изображение со стрима отправляется к нему на обработку.

Первоначально изображение отправляется на распознавание лица, если лицо распознано, только тогда изображение отправляется на распознавание маски. Каждый результат работы воркера сохраняется и может быть отображен на UI.

Скорость работы


  • Получение изображение со стрима — 31 мс
  • Препроцессинг определения лица — 0-1 мс
  • Определение лица — 51 мс
  • Постпроцессинг определения лица  — 8 мс
  • Препроцессинг определения маски — 2 мс
  • Определение маски — 11 мс
  • Постпроцессинг определения маски — 0-1 мс

Итого: 

  • Определение лица — 60 мс + 31 мс = 91 мс
  • Определение маски — 14 мс

Таким образом, за ~105 мс бы знаем всю информацию с изображения.
 
*Препроцессинг определения лица — это получение изображения со стрима и отправка в веб воркер
*Постпроцессинг определения лица — сохранение результата от воркера определения лица и его отрисовка на канвасе
*Препроцессинг определения маски — подготовка канваса с  изображением выровненного лица и передача его в веб воркер
*Постпроцессинг определения маски — сохранение результатов определения маски



Для каждой модели (определение лица и определение маски) используется отдельный веб воркер, который загружает необходимые для его работы библиотеки (OpenCV.js, Tensorflow.js, модели).

Таких воркеров у нас 3:
  • определение лица
  • определение маски
  • воркер-хелпер, который может заниматься трансформацией изображений, использовать тяжелый методы из OpenCV и Tensorflow для построения матрицы калибровки нескольких камер например.

Фичи и трюки, которые нам помогли при разработке и оптимизации


Веб воркеры и как оптимально с ними работать


Веб воркер — способ запустить скрипт в отдельном потоке.

Они  позволяют выполнять тяжелую обработку параллельно с основным потоком без блокировки пользовательского интерфейса. Основной поток выполняет логику оркестрирования, все тяжелые вычисления переносятся в веб-воркеры. Веб-воркеры поддерживаются почти во всех браузерах.



Возможности и ограничения веб-воркеров


Возможности:

  • Использование JavaScript
  • Доступ к объекту navigator
  • Доступ на чтение объекта location
  • Использование для запросов XMLHttpRequest
  • Возможность использовать setTimeout() / clearTimeout() и setInterval() / clearInterval()
  • Application Cache
  • Импорт сторонних скриптов с помощью importScripts()
  • Создание других воркеров

Ограничения:

  • Нет доступа к DOM
  • Нет доступа к объекту windows
  • Нет доступа к объекту document
  • Нет доступа к объекту parent

Общение между основным потоком и веб воркерами происходит с помощью postMessage и обработчиком событий onmessage.



Если посмотреть в спецификацию метода postMessage(), можно заметить, что он принимает не только данные, но и второй аргумент — transferable object.

worker.postMessage(message, [transfer]);

Давайте посмотрим, чем нам поможет использование его.

Transferable интерфейс представляет собой объект, который можно передавать между различными контекстами выполнения, такими как основной поток и веб-воркеры.

К ним относятся:

  • ImageBitmap 
  • OffscreenCanvas
  • ArrayBuffer
  • MessagePort

Если мы хотим передать 500 Мб данных в воркер, мы можем это сделать и без второго аргумента, но разница будет во времени передачи и использовании памяти существенная.
Передача без transfer аргумента займет 149 мс и 1042 Мб для Google Chrome, в других браузерах еще больше.



При использовании transfer это займет 1 мс и сократит потребление памяти в 2 раза!


Так как из основного потока в веб воркеры изображения передаются часто, то нам важно это делать максимально быстро и эффективно по памяти, и эта фича нам в этом очень сильно помогает.

Использование OffscreenCanvas


В веб воркере нет доступа к DOM, соответственно нельзя использовать canvas напрямую. На помощь приходит OffscreenCanvas.

Преимущества:

  • Не зависит от DOM
  • Может быть использован как в основном потоке, так и в веб воркерах
  • Имеет transferable интерфейс и не нагружает основной поток, если отрисовка происходит в веб воркере



Преимущества использования requestAnimationFrame


requestAnimationFrame позволяет получать изображения со стрима с максимальной производительностью (60 FPS) и ограничивается только возможностью камеры, не все камеры отдают видео с такой частотой.

Основными преимуществами являются:

  • Браузер оптимизирует вызовы requestAnimationFrame с другими анимациями и перерисовками, что позволяет избежать ненужных перерисовок и как следствие «лагов».
  • При использовании этого метода расход батареи значительно меньше, это особенно важно для мобильных девайсов.
  • Он работает без стека вызова, тем самым не создавая очередь вызовов.
  • Минимальная частота вызова 16.67 мс (1000 мс / 60 fps = 16.67 мс)
  • Можно контролировать частоту вызова



Снятие и анализ метрик 


Для отображения метрик приложения сейчас используется stats.js и по началу это казалось хорошей идеей, но после того, когда метрик стало 20+, основной флоу приложения начинал тормозить, из-за специфики работы браузер. Каждая метрика — это канвас, на который отрисовывается график (данные поступают очень часто туда) и браузер без остановки занимается отрисовкой, что негативно сказывается на работе приложения, следовательно и метрики заниженные получаются.

Для избежания такой проблемы лучше отказаться от использования «красоты», а выводить просто тестом значения текущее и просчитанное среднее за все время. Обновление значения в DOM будет гораздо быстрее, чем отрисовка.

Контролирование утечек памяти


Довольно часто при разработке мы сталкивались с утечкой памяти на мобильных устройствах, в то время как на десктопе работать могло очень долго.

При использовании веб воркеров нельзя узнать сколько памяти он потребляет в реальности (performance.memory не работает в веб воркерах).

На основе этого, мы предусмотрели запуск нашего приложения через веб воркеры и полностью в основном потоке. Запуская все наши модели детекции в основном потоке, можно снять метрики потребления памяти и увидеть, где утечка памяти и исправить это.

Основной код моделей в веб воркерах


Мы ознакомились с основными трюками, которые были использованы при реализации приложения, теперь рассмотрим саму реализацию.

Для работы с веб воркерами мы изначально использовали comlink-loader. Очень удобная библиотека, позволяющая работать с воркером как с объектом класса, не используя методы onmessage и postMessage, контролирование асинхронного кода с помощью async-await. Все это было удобно, пока приложение не запустили на планшете (Samsung Galaxy Tab S7) и неожиданно оно через 2 минуты работы крэшилось.

Проанализировав весь наш код, мы не нашли утечек памяти, кроме черного ящика в виде этой библиотеки для работы с воркерами. По какой-то причине запускаемые модели Tensorflow.js не очищались и где-то подвисали внутри этой библиотеки.

Было принято решение попробовать использовать worker-loader, который позволяет работать с веб воркерами как из чистого js без лишних прослоек. И это решило проблему, приложение работает сутками без вылетов.

Определение лица

Создаем воркер

this.faceDetectionWorker = workers.FaceRgbDetectionWorkerFactory.createWebWorker();

Создаем обработчик сообщений из воркера в основном потоке.

this.faceDetectionWorker.onmessage = async (event) => {
 if (event.data.type === 'load') {
   this.faceDetectionWorker.postMessage({
     type: 'init',
     backend,
     streamSettings,
     faceDetectionSettings,
     imageRatio: this.imageRatio,
   });
 } else if (event.data.type === 'init') {
   this.isFaceWorkerInit = event.data.status;

   // When both workers inited it is run processes to grab and process frames only
   if (this.isFaceWorkerInit && this.isMaskWorkerInit) {
     await this.grabFrame();
   }
 } else if (event.data.type === 'faceResults') {
   this.onFaceDetected(event);
 } else {
   throw new <i>Error</i>(`Type=${event.data.type} is not supported by RgbVideo for FaceRgbDatectionWorker`);
 }
};

Отправка изображение на обработку лица

this.faceDetectionWorker.postMessage(
 {
   type: 'detectFace',
   originalImageToProcess: this.lastImage,
   lastIndex: lastItem!.index,
 },
 [this.lastImage], // transferable object
);

Код веб воркера определения лица.

Метод init инициализирует все модели, библиотеки и канвас, которые ему пригодятся для работы.

export const init = async (data) => {
 const { backend, streamSettings, faceDetectionSettings, imageRatio } = data;

 flipHorizontal = streamSettings.flipHorizontal;
 faceMinWidth = faceDetectionSettings.faceMinWidth;
 faceMinWidthConversionFactor = faceDetectionSettings.faceMinWidthConversionFactor;
 predictionIOU = faceDetectionSettings.predictionIOU;
 recommendedLocation = faceDetectionSettings.useRecommendedLocation ? faceDetectionSettings.recommendedLocation : null;
 detectedFaceThumbnailSize = faceDetectionSettings.detectedFaceThumbnailSize;
 srcImageRatio = imageRatio;
 await tfc.setBackend(backend);
 await tfc.ready();

 const [blazeModel] = await <i>Promise</i>.all([
   blazeface.load({
     // The maximum number of faces returned by the model
     maxFaces: faceDetectionSettings.maxFaces,
     // The width of the input image
     inputWidth: faceDetectionSettings.faceDetectionImageMinWidth,
     // The height of the input image
     inputHeight: faceDetectionSettings.faceDetectionImageMinHeight,
     // The threshold for deciding whether boxes overlap too much
     iouThreshold: faceDetectionSettings.iouThreshold,
     // The threshold for deciding when to remove boxes based on score
     scoreThreshold: faceDetectionSettings.scoreThreshold,
   }),
   isOpenCvLoaded(),
 ]);

 faceDetection = new FaceDetection();
 originalImageToProcessCanvas = new <i>OffscreenCanvas</i>(srcImageRatio.videoWidth, srcImageRatio.videoHeight);
 originalImageToProcessCanvasCtx = originalImageToProcessCanvas.getContext('2d');

 resizedImageToProcessCanvas = new <i>OffscreenCanvas</i>(
   srcImageRatio.faceDetectionImageWidth,
   srcImageRatio.faceDetectionImageHeight,
 );
 resizedImageToProcessCanvasCtx = resizedImageToProcessCanvas.getContext('2d');
 return blazeModel;
};

Метод isOpenCvLoaded дожидается загрузки openCV

export const isOpenCvLoaded = () => {
 let timeoutId;

 const resolveOpenCvPromise = (resolve) => {
   if (timeoutId) {
     clearTimeout(timeoutId);
   }

   try {
     // eslint-disable-next-line no-undef
     if (cv && cv.Mat) {
       return resolve();
     } else {
       timeoutId = setTimeout(() => {
         resolveOpenCvPromise(resolve);
       }, OpenCvLoadedTimeoutInMs);
     }
   } catch {
     timeoutId = setTimeout(() => {
       resolveOpenCvPromise(resolve);
     }, OpenCvLoadedTimeoutInMs);
   }
 };

 return new <i>Promise</i>((resolve) => {
   resolveOpenCvPromise(resolve);
 });
};

Самый главный метод, это определение лица.

export const detectFace = async (data, faceModel) => {
 let { originalImageToProcess, lastIndex } = data;
 const facesThumbnailsImageData = [];

 // Resize original image to the recommended BlazeFace resolution
 resizedImageToProcessCanvasCtx.drawImage(
   originalImageToProcess,
   0,
   0,
   srcImageRatio.faceDetectionImageWidth,
   srcImageRatio.faceDetectionImageHeight,
 );
 // Getting resized image
 let resizedImageDataToProcess = resizedImageToProcessCanvasCtx.getImageData(
   0,
   0,
   srcImageRatio.faceDetectionImageWidth,
   srcImageRatio.faceDetectionImageHeight,
 );
 // Detect faces by BlazeFace
 let predictions = await faceModel.estimateFaces(
   // The image to classify. Can be a tensor, DOM element image, video, or canvas
   resizedImageDataToProcess,
   // Whether to return tensors as opposed to values
   returnTensors,
   // Whether to flip/mirror the facial keypoints horizontally. Should be true for videos that are flipped by default (e.g. webcams)
   flipHorizontal,
   // Whether to annotate bounding boxes with additional properties such as landmarks and probability. Pass in `false` for faster inference if annotations are not needed
   annotateBoxes,
 );
 // Normalize predictions
 predictions = faceDetection.normalizePredictions(
   predictions,
   returnTensors,
   annotateBoxes,
   srcImageRatio.faceDetectionImageRatio,
 );
 // Filters initial predictions by the criteri that all landmarks should be in area of interest
 predictions = faceDetection.filterPredictionsByFullLandmarks(
   predictions,
   srcImageRatio.videoWidth,
   srcImageRatio.videoHeight,
 );
 // Filters predictions by min face width
 predictions = faceDetection.filterPredictionsByMinWidth(predictions, faceMinWidth, faceMinWidthConversionFactor);
 // Filters predictions by recommended location
 predictions = faceDetection.filterPredictionsByRecommendedLocation(predictions, predictionIOU, recommendedLocation);

 // If there are any predictions it is started faces thumbnails extraction according to the configured size
 if (predictions && predictions.length > 0) {
   // Draw initial original image
   originalImageToProcessCanvasCtx.drawImage(originalImageToProcess, 0, 0);
   const originalImageDataToProcess = originalImageToProcessCanvasCtx.getImageData(
     0,
     0,
     originalImageToProcess.width,
     originalImageToProcess.height,
   );

   // eslint-disable-next-line no-undef
   let srcImageData = cv.matFromImageData(originalImageDataToProcess);
   try {
     for (let i = 0; i < predictions.length; i++) {
       const prediction = predictions[i];
       const facesOriginalLandmarks = <i>JSON</i>.parse(<i>JSON</i>.stringify(prediction.originalLandmarks));

       if (flipHorizontal) {
         for (let j = 0; j < facesOriginalLandmarks.length; j++) {
           facesOriginalLandmarks[j][0] = srcImageRatio.videoWidth - facesOriginalLandmarks[j][0];
         }
       }

       // eslint-disable-next-line no-undef
       let dstImageData = new cv.Mat();
       try {
         // eslint-disable-next-line no-undef
         let thumbnailSize = new cv.Size(detectedFaceThumbnailSize, detectedFaceThumbnailSize);

         let transformation = getOneToOneFaceTransformationByTarget(detectedFaceThumbnailSize);

         // eslint-disable-next-line no-undef
         let similarityTransformation = getSimilarityTransformation(facesOriginalLandmarks, transformation);
         // eslint-disable-next-line no-undef
         let similarityTransformationMatrix = cv.matFromArray(3, 3, cv.CV_64F, similarityTransformation.data);

         try {
           // eslint-disable-next-line no-undef
           cv.warpPerspective(
             srcImageData,
             dstImageData,
             similarityTransformationMatrix,
             thumbnailSize,
             cv.INTER_LINEAR,
             cv.BORDER_CONSTANT,
             new cv.Scalar(127, 127, 127, 255),
           );

           facesThumbnailsImageData.push(
             new <i>ImageData</i>(
               new <i>Uint8ClampedArray</i>(dstImageData.data, dstImageData.cols, dstImageData.rows),
               detectedFaceThumbnailSize,
               detectedFaceThumbnailSize,
             ),
           );
         } finally {
           similarityTransformationMatrix.delete();
           similarityTransformationMatrix = null;
         }
       } finally {
         dstImageData.delete();
         dstImageData = null;
       }
     }
   } finally {
     srcImageData.delete();
     srcImageData = null;
   }
 }

 return { resizedImageDataToProcess, predictions, facesThumbnailsImageData, lastIndex };
};

На вход подается изображение и индекс, для сопоставления лица и детекции маски в последующем.

Так как blazeface принимает изображения с максимальной стороной 128 px, то изображение с камеры нужно уменьшить.

Вызвав метод faceModel.estimateFaces мы запускаем анализ изображения с помощью  blazeface и нам возвращаются предикшены с координатами области лица, носа, ушей, глаз, рта. 
Перед тем, как с ними работать, нужно восстановить координаты для исходного изображения, мы же его сжали до 128 px.

Теперь можно использовать эти данные для принятия решения находится ли лицо в нужной области или нет, какой минимальный размер лица нам нужен, для последующей идентификации.



Следующий код вырезает лицо из изображения и выравнивает его, для идентификации и детекции маски с помощью методов openCV.

Детекция маски

Инициализация модели и webAssembly

export const init = async (data) => {
 const { backend, streamSettings, maskDetectionsSettings, imageRatio } = data;

 flipHorizontal = streamSettings.flipHorizontal;
 detectedMaskThumbnailSize = maskDetectionsSettings.detectedMaskThumbnailSize;
 srcImageRatio = imageRatio;
 await tfc.setBackend(backend);
 await tfc.ready();
 const [maskModel] = await <i>Promise</i>.all([
   tfconv.loadGraphModel(
     `/rgb_mask_classification_first/MobileNetV${maskDetectionsSettings.mobileNetVersion}_${maskDetectionsSettings.mobileNetWeight}/${maskDetectionsSettings.mobileNetType}/model.json`,
   ),
 ]);

 detectedMaskThumbnailCanvas = new <i>OffscreenCanvas</i>(detectedMaskThumbnailSize, detectedMaskThumbnailSize);
 detectedMaskThumbnailCanvasCtx = detectedMaskThumbnailCanvas.getContext('2d');
 return maskModel;
};

Для детекции маски нам необходимы координаты глаз, ушей, носа и рта и выровненное изображение, которое вернул воркер детекции лица.

this.maskDetectionWorker.postMessage({
 type: 'detectMask',
 prediction: lastItem!.data.predictions[0],
 imageDataToProcess,
 lastIndex: lastItem!.index,
});

Метод детекции

export const detectMask = async (data, maskModel) => {
 let { prediction, imageDataToProcess, lastIndex } = data;
 const masksScores = [];
 const maskLandmarks = <i>JSON</i>.parse(<i>JSON</i>.stringify(prediction.landmarks));

 if (flipHorizontal) {
   for (let j = 0; j < maskLandmarks.length; j++) {
     maskLandmarks[j][0] = srcImageRatio.faceDetectionImageWidth - maskLandmarks[j][0];
   }
 }
 // Draw thumbnail with mask
 detectedMaskThumbnailCanvasCtx.putImageData(imageDataToProcess, 0, 0);
 // Detect mask via NN
 let predictionTensor = tfc.tidy(() => {
   let maskDetectionSnapshotFromPixels = tfc.browser.<i>fromPixels</i>(detectedMaskThumbnailCanvas);
   let maskDetectionSnapshotFromPixelsFlot32 = tfc.<i>cast</i>(maskDetectionSnapshotFromPixels, 'float32');
   let expandedDims = maskDetectionSnapshotFromPixelsFlot32.expandDims(0);

   return maskModel.predict(expandedDims);
 });
 // Put mask detection result into the returned array
 try {
   masksScores.push(predictionTensor.dataSync()[0].toFixed(4));
 } finally {
   predictionTensor.dispose();
   predictionTensor = null;
 }

 return {
   masksScores,
   lastIndex,
 };
};

Результатом нейронной сети является вероятность, что маска есть, что мы и возвращаем из воркера. Это позволяет уменьшать или увеличивать трэшхолд детекции маски. По lastIndex мы можем сопоставить лицо и наличие маски и вывести на экран какую-то информацию по конкретному человеку.

Заключение


Надеюсь эта статья поможет вам узнать много нового о возможностях работы с ML в браузере и путях оптимизации. Используя описанные трюки можно оптимизировать большинство приложений.