Обнаружение лиц в Node.js с использованием Rust и WebAssembly

Будущим учащимся на курсе "Node.js Developer" предлагаем записаться на открытый урок по теме "Докеризация node.js приложений".

А пока делимся традиционным переводом статьи.

В последней статье мы рассказывали, как вызывать функции Rust из Node.js. Сегодня мы расскажем, как написать приложение AIaaS (англ. Artificial Intelligence as a Service — «искусственный интеллект как услуга») на базе Node.js.

Большинство приложений с искусственным интеллектом сейчас разрабатываются на языке Python, а главным языком программирования для веб-разработки является JavaScript. Для того чтобы реализовать возможности ИИ в вебе, нужно обернуть алгоритмы ИИ в JavaScript, а именно в Node.js.

Однако ни Python, ни JavaScript сами по себе не подходят для разработки ИИ-приложений с большим объемом вычислений. Это высокоуровневые, медленные языки со сложной средой выполнения, в которых удобство использования достигается за счет снижения производительности. Для решения этой проблемы блоки интеллектуальных вычислений в Python оборачиваются в нативные C/C++-модули. Точно так же можно сделать и в Node.js, но мы нашли решение получше — WebAssembly.

Виртуальные машины WebAssembly поддерживают тесную интеграцию с Node.js и другими средами выполнения JavaScript-кода. Они отличаются высокой производительностью, безопасны с точки зрения доступа к памяти, изначально защищены и совместимы с разными операционными системами. В нашем подходе сочетаются лучшие возможности WebAssembly и нативного кода.

Как это работает

AIaaS-приложение на базе Node.js состоит из трех компонентов.

• Приложение на Node.js является веб-сервисом и вызывает функцию WebAssembly для выполнения интенсивных вычислений, таких как алгоритмы ИИ.

• Для подготовки, постобработки и интеграции данных с другими системами используется функция WebAssembly. Мы будем использовать язык Rust. Функцию должен написать разработчик приложения.

• Для запуска модели ИИ используется нативный код, что максимально повышает быстродействие. Эта часть кода очень короткая и должна проверяться с точки зрения безопасности и защищенности. Нужно просто вызвать эту нативную программу из функции WebAssembly — точно так же, как в Python и Node.js выполняется вызов нативных функций.

А теперь рассмотрим пример.

Как можно реализовать функционал обнаружения лиц

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

Исходный код Rust для реализации модели обнаружения лиц MTCNN создан по замечательной инструкции от Cetra: Обнаружение лиц с использованием библиотеки Tensorflow на Rust. Для того чтобы Tensorflow работала в WebAssembly, мы кое-что изменили.

Приложение на Node.js отвечает за загрузку файла и вывод результата.

app.post('/infer', function (req, res) {
  let image_file = req.files.image_file;
  var result_filename = uuidv4() + ".png";

  // Call the infer() function from WebAssembly (SSVM)
  var res = infer(req.body.detection_threshold, image_file.data);

  fs.writeFileSync("public/" + result_filename, res);
  res.send('<img src="' +  result_filename + '"/>');
});

Как видите, JavaScript-приложение просто передает в функцию infer() графические данные и параметр detection_threshold, задающий минимальный размер лица, которое сможет обнаружить приложение, а затем сохраняет возвращаемое значение в графический файл на сервере. Функция infer() написана на языке Rust и скомпилирована в WebAssembly, поэтому ее можно вызвать из JavaScript.

Сначала функция infer() преобразует исходные графические данные в плоскую форму и сохраняет их в массив. Она настраивает модель TensorFlow и использует плоские графические данные в качестве входных данных. После выполнения модели TensorFlow возвращается набор чисел — координаты четырех вершин рамки, выделяющей лицо. Затем функция infer() рисует рамку вокруг каждого лица и сохраняет измененное изображение в формате PNG на веб-сервере.

#[wasm_bindgen]
pub fn infer(detection_threshold: &str, image_data: &[u8]) -> Vec<u8> {
    let mut dt = detection_threshold;
    ... ...
    let mut img = image::load_from_memory(image_data).unwrap();

    // Run the tensorflow model using the face_detection_mtcnn native wrapper
    let mut cmd = Command::new("face_detection_mtcnn");
    // Pass in some arguments
    cmd.arg(img.width().to_string())
        .arg(img.height().to_string())
        .arg(dt);
    // The image bytes data is passed in via STDIN
    for (_x, _y, rgb) in img.pixels() {
        cmd.stdin_u8(rgb[2] as u8)
            .stdin_u8(rgb[1] as u8)
            .stdin_u8(rgb[0] as u8);
    }
    let out = cmd.output();

    // Draw boxes from the result JSON array
    let line = Pixel::from_slice(&[0, 255, 0, 0]);
    let stdout_json: Value = from_str(str::from_utf8(&out.stdout).expect("[]")).unwrap();
    let stdout_vec = stdout_json.as_array().unwrap();
    for i in 0..stdout_vec.len() {
        let xy = stdout_vec[i].as_array().unwrap();
        let x1: i32 = xy[0].as_f64().unwrap() as i32;
        let y1: i32 = xy[1].as_f64().unwrap() as i32;
        let x2: i32 = xy[2].as_f64().unwrap() as i32;
        let y2: i32 = xy[3].as_f64().unwrap() as i32;
        let rect = Rect::at(x1, y1).of_size((x2 - x1) as u32, (y2 - y1) as u32);
        draw_hollow_rect_mut(&mut img, rect, *line);
    }   
    let mut buf = Vec::new();
    // Write the result image into STDOUT
    img.write_to(&mut buf, image::ImageOutputFormat::Png).expect("Unable to write");
    return buf;
}

Команда face_detection_mtcnn запускает модель TensorFlow в нативном коде. Она принимает три аргумента: ширину изображения, высоту изображения и минимальный размер лица для обнаружения. Фактические графические данные в виде плоских значений RGB передаются в программу из функции infer() WebAssembly через STDIN. Результат запуска модели кодируется в JSON и возвращается через STDOUT.

Обратите внимание, мы передали параметр модели detectiont hreshold в тензор minsize, а затем использовали тензор input для передачи входных графических данных в программу. Тензор box используется для извлечения результатов из модели.

fn main() -> Result<(), Box<dyn Error>> {
    // Get the arguments passed in from WebAssembly
    let args: Vec<String> = env::args().collect();
    let img_width: u64 = args[1].parse::<u64>().unwrap();
    let img_height: u64 = args[2].parse::<u64>().unwrap();
    let detection_threshold: f32 = args[3].parse::<f32>().unwrap();
    let mut buffer: Vec<u8> = Vec::new();
    let mut flattened: Vec<f32> = Vec::new();

    // The image bytes are read from STDIN
    io::stdin().read_to_end(&mut buffer)?;
    for num in buffer {
        flattened.push(num.into());
    }

    // Load up the graph as a byte array and create a tensorflow graph.
    let model = include_bytes!("mtcnn.pb");
    let mut graph = Graph::new();
    graph.import_graph_def(&*model, &ImportGraphDefOptions::new())?;

    let mut args = SessionRunArgs::new();
    // The `input` tensor expects BGR pixel data from the input image
    let input = Tensor::new(&[img_height, img_width, 3]).with_values(&flattened)?;
    args.add_feed(&graph.operation_by_name_required("input")?, 0, &input);

    // The `min_size` tensor takes the detection_threshold argument
    let min_size = Tensor::new(&[]).with_values(&[detection_threshold])?;
    args.add_feed(&graph.operation_by_name_required("min_size")?, 0, &min_size);

    // Default input params for the model
    let thresholds = Tensor::new(&[3]).with_values(&[0.6f32, 0.7f32, 0.7f32])?;
    args.add_feed(&graph.operation_by_name_required("thresholds")?, 0, &thresholds);
    let factor = Tensor::new(&[]).with_values(&[0.709f32])?;
    args.add_feed(&graph.operation_by_name_required("factor")?, 0, &factor);

    // Request the following outputs after the session runs.
    let bbox = args.request_fetch(&graph.operation_by_name_required("box")?, 0);

    let session = Session::new(&SessionOptions::new(), &graph)?;
    session.run(&mut args)?;

    // Get the bounding boxes
    let bbox_res: Tensor<f32> = args.fetch(bbox)?;
    let mut iter = 0;
    let mut json_vec: Vec<[f32; 4]> = Vec::new();
    while (iter * 4) < bbox_res.len() {
        json_vec.push([
            bbox_res[4 * iter + 1], // x1
            bbox_res[4 * iter],     // y1
            bbox_res[4 * iter + 3], // x2
            bbox_res[4 * iter + 2], // y2
        ]);
        iter += 1;
    }
    let json_obj = json!(json_vec);
    // Return result JSON in STDOUT
    println!("{}", json_obj.to_string()); 
    Ok(())
}

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

Как развернуть программу обнаружения лиц

Для начала нужно установить Rust, Node.js, виртуальную машину Second State WebAssembly VM и инструмент ssvmup. Ознакомьтесь с инструкцией по настройке или воспользуйтесь Docker-образом. Вам также понадобится библиотека TensorFlow.

$ wget https://storage.googleapis.com/tensorflow/libtensorflow/libtensorflow-gpu-linux-x86_64-1.15.0.tar.gz
$ sudo tar -C /usr/ -xzf libtensorflow-gpu-linux-x86_64-1.15.0.tar.gz

Для развертывания программы обнаружения лиц понадобится нативный драйвер для модели TensorFlow. Его можно скомпилировать из исходного кода на Rust, например из этого проекта.

# in the native_model_zoo/face_detection_mtcnn directory
$ cargo install --path .

Затем переходим к проекту веб-приложения. Запустите команду ssvmup для сборки функции WebAssembly из Rust.Эта функция WebAssembly осуществляет подготовку данных для веб-приложения.

# in the nodejs/face_detection_service directory
$ ssvmup build

Собрав функцию WebAssembly, можно запускать веб-приложение на Node.js.

$ npm i express express-fileupload uuid

$ cd node
$ node server.js

Теперь веб-сервис доступен на порту 8080 вашего компьютера. Попробуйте загрузить свои селфи, семейные или групповые фотографии!

TensorFlow Model Zoo

Пакет Rust face_detection_mtcnn — это обертка вокруг библиотеки TensorFlow. Она загружает обученную модель TensorFlow (так называемую замороженную сохраненную модель), передает входные данные в модель, выполняет модель и извлекает из нее выходные значения.

В данном случае наша обертка извлекает только координаты рамок вокруг обнаруженных лиц. Модель также определяет уровни достоверности для каждого обнаруженного лица и положение глаз, рта и носа. Если изменить имена тензоров для извлечения данных, обертка получит и эту информацию и вернет ее в функцию WASM.

Если вы захотите использовать другую модель, то создать для нее обертку по этому примеру будет несложно. Вам лишь нужно знать имена входного и выходного тензора и их типы данных.

Для этого мы создали проект nativemodelzoo, в рамках которого разрабатываем готовые к использованию обертки на языке Rust для самых разных моделей TensorFlow.

Что дальше

В этой статье мы разобрались, как создать AIaaS-приложение на Node.js с использованием Rust и WebAssembly для практического применения. Наш подход позволяет всем внести свой вклад в Model Zoo («зоопарк моделей»), который разработчики приложений смогут использовать в качестве ИИ-библиотеки.

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

Узнать подробнее о курсе "Node.js Developer". Зарегистрироваться на открытый урок по теме "Докеризация node.js приложений" можно здесь.