Échange de visages en temps réel avec Rust : 60 FPS sans Python et architecture sans verrous
Les solutions existantes pour l'échange de visages en temps réel sont souvent liées à Python et son écosystème, ce qui entraîne des dépendances excessives et des problèmes de performances. Cet article détaille l'implémentation d'une application en pur Rust qui atteint 60 FPS sans utiliser Python. Composants clés : ONNX Runtime pour l'inférence, des threads sans verrous, et une optimisation mémoire complète.
Pipeline de traitement des images
Chaque image passe par quatre modèles de réseaux de neurones en séquence stricte. RetinaFace détecte les visages et extrait cinq points clés. ArcFace génère un embedding 512-dimensionnel du visage source. InSwapper prend la région du visage cible et l'embedding source, renvoyant le visage échangé. GFPGAN améliore optionnellement la qualité du résultat. Tous les modèles s'exécutent via ONNX Runtime sans noyaux CUDA personnalisés — seulement des tenseurs en entrée et en sortie.
Caractéristique architecturale : pas d'allocations intermédiaires. Les tampons de pixels sont alloués une seule fois au démarrage de l'application. La conversion RGBA vers RGB, le remplissage des tenseurs et les transformations affines sont effectués sans allocations mémoire supplémentaires. La seule allocation par image est la création d'un Arc pour l'instantané final, inévitable avec ArcSwap.
Les boucles de traitement sont optimisées pour le parallélisme au niveau des instructions. Par exemple, la conversion d'espace colorimétrique traite 4 pixels par itération :
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
}
Architecture multithreading sans verrous
Le système utilise trois threads sans aucun verrou sur le chemin chaud. Le thread de capture reçoit les images via nokhwa et les publie via ArcSwap. Le thread de pipeline traite les images et passe le résultat via un second ArcSwap. Le thread UI rend le tampon actuel via egui. Optimisation critique : aligner les structures sur les lignes de cache (64 octets) pour éviter le faux partage :
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);
};
Les assertions à la compilation garantissent le respect des tailles. Les compteurs de génération atomiques remplacent les mutex, éliminant la surcharge de synchronisation. Éviter async et les channels simplifie la logique et réduit la surcharge.
Macros pour minimiser la duplication de code
Quatre réseaux de neurones nécessitent des opérations identiques : conversion RGBA vers format planaire BGR et normalisation. Au lieu de copier le code, des macros sont utilisées. La macro fill_tensor_planar! est paramétrée par le normaliseur :
// For ArcFace: (v - 127.5) / 127.5
fill_tensor_planar!(
self.tensor,
self.aligned,
PLANE,
FS,
FS,
FS * 3,
bgr_norm_arcface
);
// For InSwapper: v / 255.0
fill_tensor_planar!(
self.tensor,
self.aligned,
PLANE,
SS,
SS,
SS * 3,
bgr_inv255
);
De même, build_session! élimine l'appel triple imbriqué ort_err! lors de l'initialisation 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))?
}};
}
Transformations géométriques et UI
L'alignement des visages est implémenté avec une transformation de similarité personnalisée. Les cinq points clés de RetinaFace sont appariés au template ArcFace en résolvant un système 4x4 avec la méthode de Gauss à pivotage partiel. La transformation inverse est appliquée pour le warping bilinéaire et le collage avec fusion alpha :
let edge_dist = ax.min(ay).min(fs - ax).min(fs - ay) - margin;
let alpha = (edge_dist * inv_blend).clamp(0.0, 1.0);
Cela assure des transitions fluides aux bords sans artefacts. L'UI est construite sur egui en mode immédiat, idéal pour la vidéo en direct. Un thème personnalisé sans cadre avec des feux de circulation style macOS et un arrière-plan transparent permet de déplacer la fenêtre sans état retenu. Le rendu de chaque image sans délais maintient un 60 FPS stable.
Points clés
- Le modèle de propriété de Rust élimine les courses aux données et les use-after-free même dans une architecture multithread complexe.
- ONNX Runtime via ort s'avère prêt pour la production : chargement stable des modèles et inférence sans dépendances CUDA.
- egui pour le temps réel — rendu en mode immédiat sans état retenu minimise la latence avec les flux vidéo.
- Zéro allocations sur le chemin chaud obtenu par pré-allocation des tampons et optimisation des boucles.
- Sans verrous via ArcSwap avec alignement de cache élimine la surcharge de synchronisation inter-threads.
La solution prouve que Rust est efficace pour les tâches de vision par ordinateur haute performance. Un binaire unique sans dépendances Python s'exécute directement — téléchargez simplement la release depuis GitHub. Pour une mise à l'échelle supplémentaire des performances, passer à ONNX Runtime multi-session semble prometteur, permettant l'exécution parallèle des modèles au lieu du traitement séquentiel.
— Editorial Team
Aucun commentaire pour le moment.