Download Uformer - Uformer Download de código -fonte

Uformer

Pitão

1.0.0

Baixar

Uformer: um transformador geral em forma de U para restauração de imagens (CVPR 2022)

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

Link em papel: [arxiv] [CVPR]

Atualizar:

2024.10.24 Fix links expirados.
2022.07.06 Carregue novos códigos e modelos para o nosso Uformer.
2022.04.09 Carregue os resultados do Uformer em denoising (SIDD, DND), desverbramento de movimento (GoPro, Hide, Realblur-J/-r) e Debocus Dexuring (DPDD).
2022.03.02 O Uformer foi aceito pelo CVPR 2022!
2021.11.30 Atualize o Uformer no link ARXIV. O novo código, modelos e resultados serão enviados.
2021.10.28 Libere os resultados do Uformer32 em SIDD e DND.
2021.09.30 Libere o Uformer16 pré-treinado para Sidd Denoising.
2021.08.19 Libere um modelo pré-treinado (Uformer32)! Adicione um script para o cálculo FLOP/GMAC.
2021.07.29 Adicione um script para testar o modelo pré-treinado nas imagens de resolução arbitrária.

Neste artigo, apresentamos o Uformer, uma arquitetura baseada em transformadores eficaz e eficiente, na qual criamos uma rede hierárquica de codificadores-decodificadores usando o bloco Transformer para restauração de imagens. O Uformer possui dois projetos principais para torná -lo adequado para esta tarefa. O primeiro elemento-chave é um bloco de transformadores de janela aprimorado local, onde usamos a auto-distribuição que não é de sobreposição baseada em janelas para reduzir o requisito computacional e empregar a convolução em profundidade na rede de feed-forward para melhorar ainda mais seu potencial para capturar o contexto local. O segundo elemento-chave é que exploramos três esquemas de conexão de pular para fornecer efetivamente informações do codificador ao decodificador. Alimentado por esses dois designs, o Uformer desfruta de uma alta capacidade de capturar dependências úteis para restauração de imagens. Experiências extensas em várias tarefas de restauração de imagens demonstram a superioridade do Uformer, incluindo denoising de imagem, desbroteamento, destacando e de demonstração. Esperamos que nosso trabalho incentive mais pesquisas a explorar arquiteturas baseadas em transformadores para tarefas de visão de baixo nível.

Uformer

Dependências do pacote

O projeto é construído com Pytorch 1.9.0, Python3.7, Cuda11.1. Para dependências de pacotes, você pode instalá -las por:

pip install -r requirements.txt

Modelo pré -terenciado

Uformer_b: sidd | GoPro

Resultados do modelo pré -terenciado

Uformer_b: sidd | Dnd | GoPro | Ocultar | Realblur-j | Realblur-r | DPDD

Preparação de dados

Denoising

Para dados de treinamento do SIDD, você pode baixar o conjunto de dados SIDD-Medium a partir do URL oficial. Em seguida, gerar patches de treinamento para treinamento por:

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

Para avaliação no SIDD e DND, você pode baixar dados daqui.

Deblurrendo

Para treinamento na GoPro e avaliação no GoPro, Hide, Realblur-J e Realblur-R, você pode baixar dados daqui.

Em seguida, coloque todos os dados denoising em ../datasets/denoising e todos os dados de desverregos em ../datasets/deblurring .

Treinamento

Denoising

Para treinar o Uformer no SIDD, você pode começar o treinamento por:

sh script/train_denoise.sh

Deblurrendo

Para treinar o Uformer na GoPro, você pode começar o treinamento por:

sh script/train_motiondeblur.sh

Avaliação

Para avaliar o Uformer, você pode executar:

sh script/test.sh

Para avaliar em cada conjunto de dados, você deve descomentar a linha correspondente.

Custo computacional

Fornecemos um script simples para calcular os fracassos sozinhos, um script simples foi adicionado no model.py . Você pode alterar a configuração e executar:

 python3 model . py

O cálculo manual dos GMACs nesse repo difere um pouco do artigo principal, mas eles não influenciam a conclusão. Vamos corrigir o papel mais tarde.

Citação

Se você achar este projeto útil em sua pesquisa, considere citar:

 @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}
}