Reconstruction d’images bayésiennes utilisant des modèles génératifs profonds

Les modèles d’apprentissage automatique sont couramment entraînés de bout en bout et dans un cadre supervisé, à l’aide de données appariées (entrée, sortie). Par exemple, les méthodes récentes de super-résolution s’entraînent sur des paires d’images (faible résolution, haute résolution). Toutefois, ces approches de bout en bout nécessitent un nouvel entraînement chaque fois qu’il y a un décalage de distribution dans les entrées (par exemple, images de nuit vs images en plein jour) ou dans les variables latentes pertinentes (par exemple, flou de caméra ou mouvement de main). Dans ce travail, nous exploitons les modèles génératifs d’avant-garde (SOTA) — ici StyleGAN2 — afin de construire des prioris d’image puissants, permettant ainsi l’application du théorème de Bayes à de nombreuses tâches de reconstruction en amont. Notre méthode, appelée Reconstruction bayésienne à l’aide de modèles génératifs (BRGM), utilise un seul modèle générateur pré-entraîné pour résoudre différentes tâches de restauration d’images, telles que la super-résolution et l’imputation (in-painting), en combinant ce générateur avec différents modèles de corruption directe (forward). Nous conservons les poids du modèle générateur fixes et reconstruisons l’image en estimant l’estimation bayésienne du maximum a posteriori (MAP) du vecteur latent d’entrée qui a généré l’image reconstruite. Nous utilisons par ailleurs une inférence variationnelle pour approcher la distribution a posteriori sur les vecteurs latents, à partir de laquelle nous extrayons plusieurs solutions. Nous démontrons BRGM sur trois grands ensembles de données diversifiés : (i) 60 000 images provenant du jeu de données Flick Faces High Quality, (ii) 240 000 radiographies thoraciques issues de MIMIC-III, et (iii) une collection combinée de cinq jeux de données d’IRM cérébrale comprenant 7 329 examens. Sur ces trois ensembles de données, et sans aucun ajustement de hyperparamètres spécifique à chaque jeu, notre approche simple atteint des performances compétitives par rapport aux méthodes d’état de l’art spécifiques à chaque tâche pour la super-résolution et l’imputation, tout en étant plus généralisable et sans nécessiter d’entraînement. Le code source et les modèles pré-entraînés sont disponibles en ligne : https://razvanmarinescu.github.io/brgm/.