UniRain: Unified Image Deraining with RAG-based Dataset Distillation and Multi-objective Reweighted Optimization

Ce papier présente UniRain, un cadre unifié de dérainage d'images qui améliore la généralisation grâce à une distillation de données basée sur la génération augmentée par récupération (RAG) et une stratégie d'optimisation repondérée multi-objectifs intégrée à une architecture de mélange d'experts asymétrique, permettant ainsi de restaurer efficacement des images dégradées par la pluie et les gouttes dans diverses conditions diurnes et nocturnes.

L'essentiel

🌧️ UniRain : Le "Super-Héros" qui nettoie toutes les pluies en une seule fois

Imaginez que vous essayez de regarder un paysage magnifique à travers une vitre sale. Parfois, ce sont de longues traînées d'eau (la pluie fine), parfois ce sont de grosses gouttes qui déforment l'image, et parfois c'est la nuit, où les lumières de la ville créent un effet de flou bizarre.

Jusqu'à présent, les ordinateurs avaient besoin d'un "nettoyeur" différent pour chaque type de saleté. Un programme pour les gouttes, un autre pour la nuit, un autre pour le jour. C'était comme avoir trois balais différents dans son placard et devoir changer d'outil à chaque fois qu'on nettoyait une tache différente. C'était lent et inefficace.

UniRain, c'est le nouveau balai magique qui nettoie tout d'un coup, peu importe le type de pluie ou l'heure de la journée.

Voici comment ils ont fait, avec trois astuces principales :

1. La "Bibliothèque Intelligente" (Distillation de données par RAG)

Pour apprendre à un ordinateur à nettoyer, on lui montre des milliers d'exemples de photos sales et de photos propres. Le problème ? Les bibliothèques d'images publiques sont un peu comme un marché aux puces : il y a des trésors, mais aussi beaucoup de déchets (des images de mauvaise qualité, floues ou mal étiquetées). Si on donne tout ça à l'ordinateur, il se perd et apprend des mauvaises habitudes.

L'analogie : Imaginez que vous voulez apprendre à cuisiner un grand chef. Au lieu de lui donner 2 millions de recettes (dont 1 million sont fausses ou écrites par des enfants), vous utilisez un assistant très intelligent (basé sur l'IA générative, ou RAG).

• Cet assistant regarde chaque recette.
• Il compare la recette à des images réelles de pluie pour voir si elle est crédible.
• Il ne garde que les meilleures 2,6 % des recettes (les plus fiables).

C'est ce qu'ils appellent la "distillation de données". Ils ont filtré des millions d'images pour ne garder que le "cœur pur" de l'apprentissage. Résultat : l'ordinateur apprend plus vite et mieux, car il ne s'entraîne que sur du matériel de qualité.

2. Le "Chef d'Orchestre Équilibré" (Optimisation Multi-Objectifs)

Une fois que l'ordinateur a les bonnes images, il doit apprendre à les nettoyer. Mais il y a un problème : certaines tâches sont faciles (enlever une pluie de jour) et d'autres sont très difficiles (enlever des gouttes la nuit avec des lumières).

Si on laisse l'ordinateur apprendre seul, il va devenir un expert en tâches faciles et abandonner les tâches difficiles, car c'est plus simple pour lui. C'est comme un élève qui ne révise que les chapitres faciles pour avoir une bonne moyenne, en ignorant les maths complexes.

L'analogie : UniRain utilise un chef d'orchestre intelligent qui écoute chaque instrument (chaque type de pluie).

• Si l'instrument "Pluie de jour" joue trop fort et trop vite, le chef lui dit : "Calme-toi, on a besoin de toi, mais pas au détriment des autres".
• Si l'instrument "Pluie de nuit" a du mal, le chef lui donne plus d'attention et de temps.
• Cette technique s'appelle l'optimisation rééquilibrée. Elle force l'ordinateur à être bon partout, pas juste sur ce qui est facile.

3. L'Équipe de Spécialistes (Architecture MoE Asymétrique)

Pour traiter l'image, UniRain n'utilise pas un seul cerveau, mais une équipe d'experts.

L'analogie : Imaginez une usine de restauration d'images avec deux départements :

• Le Département d'Entrée (Encodeur) : C'est une équipe de détectives très curieux (Soft-MoE). Ils regardent l'image sale de tous les angles, en utilisant une approche douce et flexible pour comprendre tous les types de saletés possibles. Ils ne jettent rien, ils collectent tout.
• Le Département de Sortie (Décodeur) : C'est une équipe de chirurgiens très précis (Hard-MoE). Une fois qu'ils ont compris le problème, ils choisissent les meilleurs experts (les plus de 3 ou 4) pour réparer les détails fins (comme les textures des feuilles ou des bâtiments) sans toucher au reste.

En combinant ces deux approches (l'un qui explore tout, l'autre qui agit avec précision), l'image ressort nette et parfaite.

🏆 Le Résultat ?

Grâce à cette combinaison de filtres intelligents, d'équilibre parfait et d'experts spécialisés, UniRain bat tous les records actuels.

• Il nettoie aussi bien les gouttes que les traînées.
• Il fonctionne aussi bien le jour que la nuit.
• Il est même capable de nettoyer la neige et le brouillard (comme un couteau suisse de la météo).

En résumé, au lieu d'avoir un outil pour chaque type de pluie, UniRain est le couteau suisse ultime qui rend nos photos de pluie aussi claires que si le soleil avait brillé toute la journée. ☀️📸

Résumé technique

1. Problématique

Le domaine de la suppression de la pluie (deraining) a connu des progrès significatifs, mais la plupart des méthodes existantes souffrent de limitations majeures :

• Manque de généralisation : Les modèles sont généralement conçus pour des types de dégradation spécifiques (ex. : traînées de pluie diurnes, gouttes de pluie, pluie nocturne) et échouent à s'adapter à d'autres scénarios.
• Complexité du monde réel : Les scènes réelles combinent souvent plusieurs types de pluie (traînées et gouttes) dans des conditions d'éclairage variées (jour/nuit).
• Problèmes d'entraînement unifié :
  • Qualité des données : Fusionner directement tous les ensembles de données publics (synthétiques et réels, soit > 2 millions de paires) introduit un bruit important et des signaux de supervision incohérents, dégradant la performance.
  • Déséquilibre d'optimisation : Lors de l'entraînement mixte, le modèle a tendance à sur-optimiser les types de pluie « faciles » (ex. : traînées nocturnes) au détriment des types complexes (ex. : gouttes diurnes), car les taux de convergence des fonctions de perte diffèrent.

L'objectif est de développer un cadre unifié capable de restaurer des images dégradées par la pluie (traînées et gouttes, jour et nuit) en utilisant un seul modèle, tout en garantissant une robustesse et une performance élevées.

2. Méthodologie

Les auteurs proposent UniRain, un cadre unifié reposant sur trois piliers technologiques principaux :

A. Distillation de jeu de données basée sur le RAG (Retrieval-Augmented Generation)

Pour éviter les problèmes de qualité des données lors de l'entraînement mixte, les auteurs ne se contentent pas de fusionner les données brutes. Ils construisent un pipeline intelligent :

1. Phase de Récupération (Retrieval) : À partir d'un corpus massif de données publiques, le système utilise des encodeurs visuels (CLIP) et textuels (BLIP) pour créer des triplets relationnels (texte, features visuelles, image). Pour une image requête, le système effectue une recherche hiérarchique basée sur la similarité sémantique, visuelle et structurelle (SSIM) pour trouver des références réalistes.
2. Phase de Génération (Generation) : Les images requêtes et les références récupérées sont soumises à des modèles de langage-vision (VLMs) via des modèles d'ensemble (InternVL2.5, LLaVA-NeXT, MobileVLM). Ces VLMs évaluent la qualité et la fiabilité des échantillons.
3. Résultat : Un processus de vote majoritaire filtre les données, créant un ensemble « distillé » (RainRAG) contenant environ 2,6 % des données originales (52 869 paires) mais de très haute qualité, garantissant une distribution proche de la réalité.

B. Architecture Asymétrique MoE (Mixture of Experts)

Pour gérer la diversité des dégradations, UniRain utilise une architecture encodeur-décodeur asymétrique :

• Encodeur Soft-MoE : Utilise un routage « doux » (soft routing) pour combiner continuellement les contributions de plusieurs experts. Cela permet de préserver de manière collaborative les indices de dégradation divers et complexes.
• Décodeur Hard-MoE : Utilise un routage « dur » (hard routing) avec sélection Top-k. Cela active sélectivement les experts les plus pertinents pour reconstruire les textures fines et les détails structurels, optimisant ainsi l'efficacité computationnelle.

C. Optimisation Répondante Multi-objectifs (Multi-objective Reweighted Optimization)

Pour résoudre le déséquilibre d'apprentissage entre les différents types de pluie, une stratégie de ré-pondération dynamique est introduite :

• Estimation de la pente de convergence : Analyse de la dérivée de la perte (loss) pour chaque type de pluie sur une fenêtre glissante.
• Métriques adaptatives :
  • Type Balance Score (TBS) : Favorise les types de pluie qui convergent lentement pour équilibrer les contributions.
  • Type Stability Score (TSS) : Évalue la stabilité historique de la convergence pour éviter les types divergents.
  • Adaptivity Factor (AF) : Un facteur de transition qui ajuste le poids entre l'équilibre inter-types et la stabilité intra-types au fil de l'entraînement.
• Poids dynamiques : Les poids de perte sont ajustés à chaque itération pour harmoniser l'apprentissage de tous les types de pluie.

3. Contributions Clés

1. Modèle UniRain : Un cadre unifié capable de traiter simultanément les traînées et les gouttes de pluie, de jour comme de nuit, surpassant les approches spécialisées.
2. Pipeline de distillation RAG : Une méthode innovante utilisant des VLMs et la récupération augmentée pour filtrer et sélectionner automatiquement les échantillons d'entraînement les plus fiables à partir de vastes bases de données publiques.
3. Stratégie d'optimisation dynamique : Une approche de ré-pondération multi-objectifs qui résout le problème de déséquilibre de convergence dans l'entraînement mixte, assurant une performance cohérente sur tous les types de dégradation.
4. Architecture MoE Asymétrique : Une conception hybride (Soft-MoE pour l'encodeur, Hard-MoE pour le décodeur) qui équilibre expressivité et efficacité.

4. Résultats Expérimentaux

Les expériences ont été menées sur le jeu de données distillé RainRAG et plusieurs benchmarks publics réels (RealRain-1k, RainDS-real, WeatherBench).

• Performance Quantitative : UniRain obtient les meilleurs résultats moyens sur les quatre catégories (traînées/gouttes, jour/nuit). Par exemple, sur le sous-ensemble DRD (gouttes diurnes), il améliore le PSNR de 1,35 dB par rapport à Restormer et 1,41 dB par rapport à NeRD-Rain.
• Performance sur Données Réelles : Sur les benchmarks réels, UniRain dépasse l'état de l'art (URIR) d'une moyenne de 1,73 dB en PSNR. Il parvient même à supprimer des gouttes résiduelles présentes dans les images de référence (Ground Truth), produisant des résultats visuellement supérieurs.
• Généralisation : Le modèle démontre une forte capacité de généralisation sur des scénarios variés (conduite autonome, drones, maritime) sans réentraînement.
• Efficacité : Malgré sa complexité, UniRain présente un nombre de paramètres et de FLOPs compétitif (voire inférieur) par rapport à des modèles lourds comme Restormer ou DRSformer.
• Extensions : Le cadre a également été étendu avec succès à la restauration météo « tout-en-un » (pluie, neige, brouillard), surpassant les méthodes dédiées à la restauration météo.

5. Signification et Impact

Ce travail est significatif car il s'attaque au problème fondamental de la fragmentation dans le domaine de la restauration d'images. Au lieu de développer des modèles spécifiques pour chaque condition, UniRain propose une approche unifiée et robuste.

• Innovation Data-Centric : L'utilisation du RAG pour la distillation de données ouvre une nouvelle voie pour améliorer la qualité des ensembles de données d'entraînement sans nécessiter de nouvelles captures coûteuses.
• Robustesse Opérationnelle : En résolvant les déséquilibres d'optimisation, le modèle devient plus fiable pour des applications réelles où les conditions de pluie sont imprévisibles et mélangées.
• État de l'art : Les résultats démontrent qu'un modèle unifié peut surpasser les modèles spécialisés, établissant une nouvelle référence pour la suppression de la pluie dans des environnements complexes.