Téléchargement Uformer - Uformer le code source Téléchargement

Uformer

Python

1.0.0

Télécharger

UFormer: un transformateur général en forme de U pour la restauration d'images (CVPR 2022)

Zhendong Wang, Xiaodong Cun, Jianmin Bao, Wengang Zhou, Jianzhuang Liu, Houqiang Li

Lien papier: [arXiv] [CVPR]

Mise à jour:

2024.10.24 Correction des liens expirés.
2022.07.06 Téléchargez de nouveaux codes et modèles pour notre uforger.
2022.04.09 Téléchargez les résultats de UFormer sur Denooing (SIDD, DND), Motion DeBurring (GoPro, Hide, RealBlur-J / -R) et Defocus Deblurring (DPDD).
2022.03.02 Uformrer a été accepté par CVPR 2022!
2021.11.30 Mise à jour Uformrer dans le lien arXiv. Le nouveau code, les modèles et les résultats seront téléchargés.
2021.10.28 Libérez les résultats de UFormer32 sur SIDD et DND.
2021.09.30 Libération pré-formée UFormer16 pour SIDD Denoising.
2021.08.19 Libérez un modèle pré-formé (UFormer32)! Ajoutez un script pour le calcul Flop / GMAC.
2021.07.29 Ajoutez un script pour tester le modèle pré-formé sur les images de résolution arbitraire.

Dans cet article, nous présentons UFormer, une architecture basée sur un transformateur efficace et efficiente, dans laquelle nous construisons un réseau hiérarchique-codeur-décodeur en utilisant le bloc de transformateur pour la restauration d'images. UFormer a deux conceptions de base pour la rendre adaptée à cette tâche. Le premier élément clé est un bloc de transformateur de fenêtres à amélioration locale, où nous utilisons une auto-atténuation de fenêtre non chevauchante pour réduire les besoins de calcul et utiliser la convolution en profondeur dans le réseau Feed-forward pour améliorer encore son potentiel de capture de contexte local. Le deuxième élément clé est que nous explorons trois schémas de connexion à saut pour fournir efficacement les informations de l'encodeur au décodeur. Propulsé par ces deux conceptions, UFormer bénéficie d'une grande capacité pour capturer des dépendances utiles pour la restauration d'images. Des expériences étendues sur plusieurs tâches de restauration d'images démontrent la supériorité de l'uFormer, y compris le débroussage d'image, la déraisonment, le désblurage et le démontage. Nous nous attendons à ce que notre travail encourage de nouvelles recherches pour explorer les architectures basées sur les transformateurs pour les tâches de vision de bas niveau.

Uforger

Dépendances de package

Le projet est construit avec Pytorch 1.9.0, Python3.7, CUDA11.1. Pour les dépendances de package, vous pouvez les installer par:

pip install -r requirements.txt

Modèle pré-entraîné

UFormer_B: SIDD | Bopro

Résultats du modèle pré-entraîné

UFormer_B: SIDD | MND | GoPro | Masquer | Realblur-j | Realblur-r | Dpdd

Préparation des données

Émoi

Pour la formation des données de SIDD, vous pouvez télécharger l'ensemble de données SIDD-Medium à partir de l'URL officielle. Générez ensuite des patchs de formation pour la formation en:

 python3 generate_patches_SIDD . py - - src_dir .. / SIDD_Medium_Srgb / Data - - tar_dir .. / datasets / denoising / sidd / train

Pour l'évaluation sur SIDD et DND, vous pouvez télécharger des données à partir d'ici.

Déblure

Pour la formation sur GoPro et l'évaluation sur GoPro, Hide, RealBlur-J et Realblur-R, vous pouvez télécharger des données à partir d'ici.

Ensuite, mettez toutes les données de dénaison dans ../datasets/denoising , et toutes les données de déblure dans ../datasets/deblurring .

Entraînement

Émoi

Pour former UFormer sur SIDD, vous pouvez commencer la formation en:

sh script/train_denoise.sh

Déblure

Pour former UFormer sur GoPro, vous pouvez commencer la formation en:

sh script/train_motiondeblur.sh

Évaluation

Pour évaluer UFormer, vous pouvez exécuter:

sh script/test.sh

Pour évaluer sur chaque ensemble de données, vous devez décommenter la ligne correspondante.

Coût de calcul

Nous fournissons un script simple pour calculer les flops par nous-mêmes, un simple script a été ajouté dans model.py . Vous pouvez modifier la configuration et l'exécution:

 python3 model . py

Le calcul manuel des GMAC dans ce repo diffère légèrement de l'article principal, mais ils n'influencent pas la conclusion. Nous corrigerons le papier plus tard.

Citation

Si vous trouvez ce projet utile dans vos recherches, veuillez envisager de citer:

 @InProceedings{Wang_2022_CVPR,
    author    = {Wang, Zhendong and Cun, Xiaodong and Bao, Jianmin and Zhou, Wengang and Liu, Jianzhuang and Li, Houqiang},
    title     = {Uformer: A General U-Shaped Transformer for Image Restoration},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    month     = {June},
    year      = {2022},
    pages     = {17683-17693}
}