Página del proyecto | Papel | ¿Modelo de tarjeta?
Nuestro trabajo de seguimiento Restauración de imágenes fotorrealistas en la naturaleza con modelos de visión y lenguaje controlados (CVPRW 2024) presenta un muestreo posterior para una mejor generación de imágenes y maneja imágenes de degradación mixta del mundo real similares a Real-ESRGAN.
[ 2024.04.16 ] ¡Nuestro artículo de seguimiento "Restauración de imágenes fotorrealistas en la naturaleza con modelos de visión y lenguaje controlados" ya está en ArXiv!
[ 2024.04.15 ] Se actualizó un modelo de IR salvaje para degradaciones del mundo real y el muestreo posterior para una mejor generación de imágenes. Los pesos previamente entrenados wild-ir.pth y wild-daclip_ViT-L-14.pt también se proporcionan para wild-ir.
[ 2024.01.20 ] ??? ¿Nuestro artículo DA-CLIP fue aceptado por ICLR 2024? Además, proporcionamos un modelo más robusto en la tarjeta modelo.
[ 2023.10.25 ] Se agregaron enlaces de conjuntos de datos para capacitación y pruebas.
[ 2023.10.13 ] Se agregó la demostración y la API de Replicate. Gracias a @chenxwh!!! Actualizamos la demostración de Hugging Face y la demostración de Colab en línea. Gracias a @fffiloni y @camenduru!!! ¿También hicimos una tarjeta modelo en Hugging Face? y proporcionó más ejemplos para realizar pruebas.
[ 2023.10.09 ] Los pesos previamente entrenados de DA-CLIP y el modelo Universal IR se publican en link1 y link2, respectivamente. Además, también proporcionamos un archivo de aplicación Gradio en caso de que desee probar sus propias imágenes.
Te aconsejamos primero crear un entorno virtual con:
python3 -m venv .env
source .env/bin/activate
pip install -U pip
pip install -r requirements.txt
Ingrese al directorio universal-image-restoration y ejecute:
import torch
from PIL import Image
import open_clip
checkpoint = 'pretrained/daclip_ViT-B-32.pt'
model , preprocess = open_clip . create_model_from_pretrained ( 'daclip_ViT-B-32' , pretrained = checkpoint )
tokenizer = open_clip . get_tokenizer ( 'ViT-B-32' )
image = preprocess ( Image . open ( "haze_01.png" )). unsqueeze ( 0 )
degradations = [ 'motion-blurry' , 'hazy' , 'jpeg-compressed' , 'low-light' , 'noisy' , 'raindrop' , 'rainy' , 'shadowed' , 'snowy' , 'uncompleted' ]
text = tokenizer ( degradations )
with torch . no_grad (), torch . cuda . amp . autocast ():
text_features = model . encode_text ( text )
image_features , degra_features = model . encode_image ( image , control = True )
degra_features /= degra_features . norm ( dim = - 1 , keepdim = True )
text_features /= text_features . norm ( dim = - 1 , keepdim = True )
text_probs = ( 100.0 * degra_features @ text_features . T ). softmax ( dim = - 1 )
index = torch . argmax ( text_probs [ 0 ])
print ( f"Task: { task_name } : { degradations [ index ] } - { text_probs [ 0 ][ index ] } " )Preparar los conjuntos de datos de entrenamiento y prueba siguiendo nuestra sección de Construcción de conjuntos de datos en papel como:
# ### for training dataset ####
# ### (uncompleted means inpainting) ####
datasets/universal/train
| --motion-blurry
| | --LQ/ * .png
| | --GT/ * .png
| --hazy
| --jpeg-compressed
| --low-light
| --noisy
| --raindrop
| --rainy
| --shadowed
| --snowy
| --uncompleted
# ### for testing dataset ####
# ### (the same structure as train) ####
datasets/universal/val
...
# ### for clean captions ####
datasets/universal/daclip_train.csv
datasets/universal/daclip_val.csv Luego ingrese al directorio universal-image-restoration/config/daclip-sde y modifique las rutas del conjunto de datos en los archivos de opciones en options/train.yml y options/test.yml .
Puede agregar más tareas o conjuntos de datos a los directorios train y val y agregar la palabra degradación a distortion .
| Degradación | movimiento borroso | brumoso | jpeg comprimido* | poca luz | ruidoso* (igual que jpeg) |
|---|---|---|---|---|---|
| Conjuntos de datos | gopro | RESIDENCIA-6k | DIV2K+Flickr2K | JAJAJA | DIV2K+Flickr2K |
| Degradación | gota de agua | lluvioso | sombreado | de mucha nieve | incompleto |
|---|---|---|---|---|---|
| Conjuntos de datos | Gota de agua | Rain100H: entrenar, probar | SDR | Nieve100K | CelebaHQ-256 |
Solo debe extraer los conjuntos de datos del tren para el entrenamiento , y todos los conjuntos de datos de validación se pueden descargar en Google Drive. Para conjuntos de datos jpeg y ruidosos, puede generar imágenes LQ utilizando este script.
Consulte DA-CLIP.md para obtener más detalles.
El código principal para el entrenamiento está en universal-image-restoration/config/daclip-sde y la red central para DA-CLIP está en universal-image-restoration/open_clip/daclip_model.py .
Coloque los pesos DA-CLIP previamente entrenados en el directorio pretrained y verifique la ruta daclip .
Luego puede entrenar el modelo siguiendo los siguientes scripts de bash:
cd universal-image-restoration/config/daclip-sde
# For single GPU:
python3 train.py -opt=options/train.yml
# For distributed training, need to change the gpu_ids in option file
python3 -m torch.distributed.launch --nproc_per_node=2 --master_port=4321 train.py -opt=options/train.yml --launcher pytorch Los modelos y registros de entrenamiento se guardarán en log/universal-ir . Puede imprimir su registro en ese momento ejecutando tail -f log/universal-ir/train_universal-ir_***.log -n 100 .
Se pueden utilizar los mismos pasos de capacitación para la restauración de imágenes en estado salvaje (wild-ir).
| Nombre del modelo | Descripción | Google Drive | AbrazosCara |
|---|---|---|---|
| DA-CLIP | Modelo CLIP consciente de la degradación | descargar | descargar |
| IR universal | Modelo de restauración de imágenes universal basado en DA-CLIP | descargar | descargar |
| Mezcla DA-CLIP | Modelo CLIP sensible a la degradación (agregue desenfoque gaussiano + pintura de rostros y desenfoque gaussiano + Rainy) | descargar | descargar |
| Mezcla IR universal | Modelo de restauración de imágenes universal basado en DA-CLIP (agregue entrenamiento sólido y degradaciones mixtas) | descargar | descargar |
| Salvaje-DA-CLIP | Modelo CLIP consciente de la degradación en estado salvaje (ViT-L-14) | descargar | descargar |
| IR salvaje | Modelo de restauración de imágenes basado en DA-CLIP en la naturaleza | descargar | descargar |
Para evaluar nuestro método de restauración de imágenes, modifique la ruta de referencia y la ruta del modelo y ejecute
cd universal-image-restoration/config/universal-ir
python test.py -opt=options/test.yml Aquí proporcionamos un archivo app.py para probar sus propias imágenes. Antes de eso, debe descargar los pesos previamente entrenados (DA-CLIP y UIR) y modificar la ruta del modelo en options/test.yml . Luego, simplemente ejecutando python app.py , puedes abrir http://localhost:7860 para probar el modelo. (También proporcionamos varias imágenes con diferentes degradaciones en el directorio images ). También proporcionamos más ejemplos de nuestro conjunto de datos de prueba en Google Drive.
Se pueden utilizar los mismos pasos para la restauración de imágenes en estado salvaje (wild-ir).
? En las pruebas, descubrimos que el modelo actual preentrenado todavía tiene dificultades para procesar algunas imágenes del mundo real que podrían tener cambios de distribución con nuestro conjunto de datos de entrenamiento (capturadas desde diferentes dispositivos o con diferentes resoluciones o degradaciones). ¡Lo consideramos como un trabajo futuro e intentaremos hacer nuestro modelo más práctico! También animamos a los usuarios interesados en nuestro trabajo a entrenar sus propios modelos con un conjunto de datos más grande y más tipos de degradación.
? Por cierto, también descubrimos que cambiar el tamaño de las imágenes de entrada directamente provocará un rendimiento deficiente en la mayoría de las tareas . Podríamos intentar agregar el paso de cambio de tamaño al entrenamiento, pero siempre destruye la calidad de la imagen debido a la interpolación.
? Para la tarea de pintura interna, nuestro modelo actual solo admite pintura facial debido a la limitación del conjunto de datos. Proporcionamos nuestros ejemplos de máscaras y puede utilizar el script generate_masked_face para generar caras incompletas.
Reconocimiento: Nuestro DA-CLIP se basa en IR-SDE y open_clip. ¡Gracias por su código!
Si tiene alguna pregunta, comuníquese con: [email protected]
Si nuestro código le ayuda en su investigación o trabajo, considere citar nuestro artículo. Las siguientes son referencias de BibTeX:
@article{luo2023controlling,
title={Controlling Vision-Language Models for Universal Image Restoration},
author={Luo, Ziwei and Gustafsson, Fredrik K and Zhao, Zheng and Sj{"o}lund, Jens and Sch{"o}n, Thomas B},
journal={arXiv preprint arXiv:2310.01018},
year={2023}
}
@article{luo2024photo,
title={Photo-Realistic Image Restoration in the Wild with Controlled Vision-Language Models},
author={Luo, Ziwei and Gustafsson, Fredrik K and Zhao, Zheng and Sj{"o}lund, Jens and Sch{"o}n, Thomas B},
journal={arXiv preprint arXiv:2404.09732},
year={2024}
}
