Face swap na Rustu: 60 FPS bez Pythonu | Technický rozbor

# Výměna obličejů v reálném čase na Rustu: 60 FPS bez Pythonu a lock-free architektura

Stávající řešení pro výměnu obličejů v reálném čase jsou často svázána s Pythonem a jeho ekosystémem, což vede k zbytečným závislostem a problémům s výkonem. V tomto článku rozebereme implementaci aplikace v čistém Rustu, která dosahuje 60 snímků za sekundu bez použití Pythonu. Klíčové komponenty: ONNX Runtime pro inferenci, lock-free vlákna a plná optimalizace paměti.

Řetězec zpracování snímků

Každý snímek prochází čtyřmi neuronovými sítěmi v přísném pořadí. RetinaFace detekuje obličeje a extrahuje pět klíčových bodů. ArcFace generuje 512-rozměrný embedding zdrojového obličeje. InSwapper přijímá oblast cílového obličeje a embedding zdroje, vrací vyměněný obličej. GFPGAN volitelně zlepšuje kvalitu výsledku. Všechny modely pracují přes ONNX Runtime bez vlastních CUDA jader — pouze tenzory na vstupu a výstupu.

Peculiarita architektury: absence meziúložištních alokací. Buffery pixelů se alokují jednou při startu aplikace. Konverze RGBA do RGB, vyplňování tenzorů a affinní transformace se provádějí bez dalších alokací paměti. Jediná alokace na snímek — vytvoření Arc pro finální snapshot, což je nevyhnutelné při použití ArcSwap.

Smyčky zpracování jsou optimalizovány pro paralelizmus na úrovni instrukcí. Například konverze barevného prostoru zpracovává 4 pixely za iteraci:

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];
    // ... ještě 2 pixely
}

Architektura vícevláknovosti bez zámků

Systém využívá tři vlákna s nulovými zámky na horké cestě. Vlákno zachycení získává snímky přes nokhwa a publikuje je přes ArcSwap. Vlákno řetězce zpracovává snímky a předává výsledek přes druhý ArcSwap. UI vlákno renderuje aktuální buffer přes egui. Kritická optimalizace — zarovnání struktur podle cache line (64 bajtů) pro prevenci false sharing:

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);
};

Compile-time asserty zaručují dodržení velikosti. Atomové čítače generací nahrazují mutexy, vylučují overhead synchronizace. Odmítnutí async a kanálů zjednodušuje logiku a snižuje režii.

Makra pro minimalizaci duplikace kódu

Čtyři neuronové sítě vyžadují stejné operace: konverze RGBA do BGR planar formátu, normalizace. Místo kopírování kódu jsou použita makra. Makro fill_tensor_planar! je parametrizováno normalizátorem:

// Pro ArcFace: (v - 127.5) / 127.5
fill_tensor_planar!(
    self.tensor,
    self.aligned,
    PLANE,
    FS,
    FS,
    FS * 3,
    bgr_norm_arcface
);

// Pro InSwapper: v / 255.0
fill_tensor_planar!(
    self.tensor,
    self.aligned,
    PLANE,
    SS,
    SS,
    SS * 3,
    bgr_inv255
);

Analogně build_session! eliminuje trojnásobné vnořené volání ort_err! při inicializaci 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))?
    }};
}

Geometrické transformace a UI

Zarovnání obličejů je realizováno vlastní similarity transformací. Pět klíčových bodů RetinaFace se mapuje na šablonu ArcFace řešením systému 4x4 metodou Gauss-Seidel s částečným pivotováním. Inverzní transformace se aplikuje pro bilinear warp a paste-back s alfa-blendingem:

let edge_dist = ax.min(ay).min(fs - ax).min(fs - ay) - margin;
let alpha = (edge_dist * inv_blend).clamp(0.0, 1.0);

To zajišťuje plynulý přechod po okrajích bez artefaktů. UI je postaveno na egui v immediate mode, což je ideální pro live video. Vlastní frameless téma s macOS-style traffic lights a průhledným pozadím umožňuje táhnutí okna bez retained state. Rendering každého snímku bez zpoždění podporuje stabilních 60 FPS.

Co je důležité

• Ownership model Rustu vyloučil data race a use-after-free i v složité vícevláknové architektuře.
• ONNX Runtime přes ort se ukázal být připravený pro produkci: stabilní načítání modelů a inferenci bez CUDA závislostí.
• egui pro real-time — immediate mode rendering bez retained state minimalizuje zpoždění při práci s videopotokem.
• Nula alokací na horké cestě dosažena předem alokovanými buffery a optimalizací smyček.
• Lock-free přes ArcSwap s cache zarovnáním eliminovalo overhead synchronizace mezi vlákny.

Řešení dokazuje, že Rust je efektivní pro high-performance úlohy počítačového vidění. Jediný binárník bez Python závislostí se spouští „z krabice“ — stačí stáhnout release z GitHubu. Pro další růst výkonu je perspektivní přechod na multisession ONNX Runtime, umožňující paralelní spouštění modelů místo sekvenčního zpracování.

— Editorial Team