Большинство инструментов для замены лиц — это Python-скрипты, склеенные из PyTorch, OpenCV и надежды. Они работают, но тащат за собой гигабайты зависимостей, требуют правильно настроенного CUDA и разваливаются в тот момент, когда ты пытаешься запустить их в реальном времени.
Мне стало интересно: можно ли собрать весь пайплайн на чистом Rust? Без Python. Без PyTorch. Без обёрток. Один бинарник, который скачал, распаковал и запустил.
Оказалось, можно. 60 fps на веб-камере.
Пайплайн
На каждом кадре последовательно отрабатывают четыре нейросети.
RetinaFace находит лица и извлекает пять ключевых точек. ArcFace вычисляет 512-мерный эмбеддинг исходного лица. InSwapper принимает регион целевого лица и эмбеддинг источника, на выходе отдаёт заменённое лицо. GFPGAN опционально улучшает результат для более высокого качества.
Все четыре модели работают через ONNX Runtime. Никаких кастомных CUDA-ядер, никакого оверхеда фреймворков. Тензор на вход, тензор на выход.
Архитектура потоков
Три потока, ноль блокировок на горячем пути.
Поток захвата получает кадры с веб-камеры через nokhwa и публикует их через ArcSwap. Поток пайплайна подхватывает новые кадры, прогоняет инференс и публикует обработанные кадры через второй ArcSwap. Поток UI читает актуальный буфер и рендерит через egui.
Никаких мьютексов на данных кадра. Никаких каналов. Никакого async. Только атомарные счётчики поколений и lock-free замена указателей.
Структуры разделяемого состояния занимают ровно по 64 байта каждая и выровнены по кэш-линиям, чтобы исключить false sharing между ядрами. Это проверяется compile-time ассертами:
const CL: usize = 64; #[repr(C, align(64))] struct ProducerCell { frame: ArcSwap<FrameSnapshot>, generation: AtomicU64, pad: [u8; CL - 16], } const : () = { assert!(std::mem::size_of::<ProducerCell>() == CL); };
Ноль аллокаций на горячем пути
Все пиксельные буферы в пайплайне аллоцируются один раз при старте. Конверсия RGBA в RGB, заполнение тензоров, аффинный варп, блендинг paste-back — ни одна из этих операций не выделяет память в процессе обработки. Единственная аллокация на кадр — Arc для финального снапшота, это неизбежно при паттерне с ArcSwap.
Конверсия RGBA в RGB развёрнута на 4 пикселя за итерацию для лучшего instruction-level parallelism:
for i in 0..chunks { let si = i 16; let di = i 12; rgb[di] = rgba[si]; rgb[di + 1] = rgba[si + 1]; rgb[di + 2] = rgba[si + 2]; rgb[di + 3] = rgba[si + 4]; rgb[di + 4] = rgba[si + 5]; rgb[di + 5] = rgba[si + 6]; // ... ещё 2 пикселя }
Макросы вместо дублирования
Каждая из четырёх нейросетей делает одно и то же: читает interleaved RGB, перекладывает в BGR planar формат, применяет нормализацию. Отличается только скалярное преобразование. Вместо четырёх копий одного и того же цикла используется макрос fill_tensor_planar!, который параметризуется нормализатором:
// Для распознавателя: (v - 127.5) / 127.5 fill_tensor_planar!( self.tensor, self.aligned, PLANE, FS, FS, FS 3, bgr_norm_arcface ); // Для свопера: v / 255.0 fill_tensor_planar!( self.tensor, self.aligned, PLANE, SS, SS, SS 3, bgr_inv255 );
Аналогично build_session! убирает тройной вложенный вызов ort_err! при создании сессии ONNX Runtime, который иначе копировался в четырёх конструкторах:
macro_rules! build_session { ($path:expr, $threads:expr) => {{ let builder = ort_err!(Session::builder())?; let mut builder = ort_err!(builder.with_intra_threads($threads))?; ort_err!(builder.commit_from_file($path))? }}; }
Аффинные преобразования
Для выравнивания лиц используется собственная реализация similarity transform. Пять ключевых точек от RetinaFace сопоставляются с шаблоном ArcFace через решение системы 4×4 методом Гаусса с частичным выбором ведущего элемента. Обратное преобразование используется для bilинейного варпа и финального paste-back с альфа-блендингом по расстоянию до края:
let edge_dist = ax.min(ay).min(fs - ax).min(fs - ay) - margin; let alpha = (edge_dist * inv_blend).clamp(0.0, 1.0);
Это даёт плавный переход по краям без видимых артефактов.
UI
Кастомная frameless тема с macOS-style traffic lights, прозрачным фоном и drag-to-move. egui в immediate mode идеально подходит для live-видео: никакого retained state, текстура обновляется каждый кадр, рендер не создаёт задержек при 60fps.
Что я вынес
Rust реально хорош для таких задач. Ownership модель сделала многопоточную архитектуру тривиальной. Никаких data race, никаких use-after-free, никаких загадочных крэшей в три часа ночи.
ONNX Runtime через крейт ort готов к продакшену. Загрузка моделей, создание тензоров, инференс — всё работает.
egui недооценён для real-time приложений. Immediate mode рендеринг без retained state идеально подходит для live-видео.
Попробовать
В архиве релиза лежит бинарник и все модели. Скачал, распаковал, запустил. Больше ничего не нужно.
Буду рад обратной связи по архитектуре и коду. Особенно интересно, есть ли у кого-то опыт с мультисессионным ONNX Runtime — запуск нескольких моделей параллельно вместо последовательного мог бы выжать ещё больше fps.
ссылка на оригинал статьи https://habr.com/ru/articles/1024700/