Face swap na Rust: 60 FPS bez Pythona | Analiza techniczna

# Zamiana twarzy w czasie rzeczywistym na Rust: 60 FPS bez Pythona i architektura bez blokad

Istniejące rozwiązania do zamiany twarzy w czasie rzeczywistym często opierają się na Pythonie i jego ekosystemie, co prowadzi do nadmiernych zależności i problemów z wydajnością. W tym artykule omawiamy implementację aplikacji w czystym Rust, osiągającej 60 klatek na sekundę bez użycia Pythona. Kluczowe komponenty: ONNX Runtime do inferencji, strumienie bez blokad i pełna optymalizacja pamięci.

Potok przetwarzania klatek

Każda klatka przechodzi przez cztery modele neuronowe w ścisłej kolejności. RetinaFace wykrywa twarze i wyciąga pięć kluczowych punktów. ArcFace generuje 512-wymiarowe osadzenie źródłowej twarzy. InSwapper przyjmuje region docelowej twarzy i osadzenie źródła, zwracając zastąpioną twarz. GFPGAN opcjonalnie poprawia jakość wyniku. Wszystkie modele działają za pośrednictwem ONNX Runtime bez niestandardowych jąder CUDA — tylko tensory na wejściu i wyjściu.

Cechą architektury jest brak pośrednich alokacji. Bufory pikseli są alokowane jednorazowo przy starcie aplikacji. Konwersja RGBA na RGB, wypełnianie tensorów i transformacje afiniczne są wykonywane bez dodatkowych alokacji pamięci. Jedyna alokacja na klatkę to utworzenie Arc dla ostatecznego snapshotu, co jest nieuniknione przy użyciu ArcSwap.

Pętle przetwarzania są zoptymalizowane pod kątem paralelizmu na poziomie instrukcji. Na przykład konwersja przestrzeni kolorów przetwarza 4 piksele na iterację:

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];
    // ... also 2 pikselya
}

Architektura wielowątkowości bez blokad

System wykorzystuje trzy wątki z zerowymi blokadami na gorącej ścieżce. Wątek przechwytywania pobiera klatki za pośrednictwem nokhwa i publikuje je przez ArcSwap. Wątek potoku przetwarza klatki i przekazuje wynik przez drugi ArcSwap. Wątek UI renderuje aktualny bufor za pomocą egui. Kluczowa optymalizacja to wyrównanie struktur do linii cache (64 bajty), aby zapobiec 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);
};

Asercje w czasie kompilacji gwarantują zachowanie rozmiaru. Atomowe liczniki generacji zastępują muteksy, eliminując narzut synchronizacji. Rezygnacja z async i kanałów upraszcza logikę i zmniejsza koszty.

Makra do minimalizacji duplikacji kodu

Cztery sieci neuronowe wymagają identycznych operacji: konwersja RGBA na format BGR planar, normalizacja. Zamiast kopiowania kodu użyto makr. Makro fill_tensor_planar! jest parametryzowane normalizatorem:

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

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

Podobnie build_session! eliminuje potrójne zagnieżdżone wywołanie ort_err! przy inicjalizacji 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))?
    }};
}

Transformacje geometryczne i UI

Wyrównanie twarzy jest realizowane za pomocą własnej transformacji podobieństwa. Pięć kluczowych punktów RetinaFace jest mapowanych na szablon ArcFace przez rozwiązanie układu 4x4 metodą Gaussa z częściowym wyborem elementu głównego. Odwrotna transformacja jest stosowana do bilinearnego warp i paste-back z alpha-blendingiem:

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 zapewnia płynne przejście na krawędziach bez artefaktów. UI jest zbudowane na egui w trybie immediate, co jest idealne dla wideo na żywo. Niestandardowa tema bez ramki z macOS-style traffic lights i przezroczystym tłem pozwala na przeciąganie okna bez retained state. Renderowanie każdej klatki bez opóźnień utrzymuje stabilne 60 FPS.

Co jest ważne

• Model własności w Rust wykluczył data race i use-after-free nawet w złożonej architekturze wielowątkowej.
• ONNX Runtime przez ort okazał się gotowy do produkcji: stabilne ładowanie modeli i inferencja bez zależności CUDA.
• egui do real-time — renderowanie w trybie immediate bez retained state minimalizuje opóźnienia przy pracy z strumieniem wideo.
• Zero alokacji na gorącej ścieżce osiągnięto przez wstępne alokowanie buforów i optymalizację pętli.
• Bez blokad przez ArcSwap z wyrównaniem cache eliminuje narzut synchronizacji między wątkami.

Rozwiązanie dowodzi, że Rust jest efektywny w zadaniach high-performance computer vision. Pojedynczy binarny bez zależności Pythona uruchamia się „z pudełka” — wystarczy pobrać release z GitHub. Dla dalszego wzrostu wydajności obiecujący jest przejście na wielosesyjny ONNX Runtime, umożliwiający równoległe uruchamianie modeli zamiast sekwencyjnego przetwarzania.

— Editorial Team