If you are interested in image super-resolution reconstruction, then you have likely heard of SwinIR. SwinIR is an image restoration model based on the Swin Transformer that achieves state-of-the-art performance in tasks such as classical/lightweight/practical image super-resolution, grayscale/color image denoising, and JPEG compression artifact removal. In this article, we will introduce how to use SwinIR for image super-resolution reconstruction.



1. Purchase a Graphics Card

I highly recommend purchasing a graphics card before you start. The Nvidia GeForce RTX 4090 is a great choice.


After assembling your machine, install the operating system and graphics card driver. For easier debugging, I recommend using a Linux system.

You must install Nvidia's specialized Docker, which supports passing the graphics card through to the Docker container. Detailed tutorials can be found here.

Complete the basic Nvidia environment and Docker environment setup, ensuring that the GPU can be used in Docker, with CUDA and cuDNN properly configured.

2. Configure the Basic Container

First, write a basic Dockerfile with Pytorch and related dependencies.

FROM nvidia/cuda:11.8.0-devel-ubuntu22.04

RUN apt update
RUN apt upgrade -y

ENV DEBIAN_FRONTEND=noninteractive

RUN DEBIAN_FRONTEND=noninteractive apt-get update && \
    apt-get install -y tzdata && \
    echo "Etc/UTC" > /etc/timezone && \
    ln -fs /usr/share/zoneinfo/UTC /etc/localtime && \
    dpkg-reconfigure -f noninteractive tzdata

RUN apt install -y nvidia-cuda-toolkit \
    gcc g++ make cmake git sudo wget curl vim python3 python3-pip python-is-python3 \
    libgl1-mesa-glx libglib2.0-0 libsm6 libxext6 libxrender-dev libxext-dev libxrender-dev

This container will be used for many different AI applications. So, I compiled it into a base image: hub.aiursoft.cn/aiursoft/internalimages/nvidia:latest.

You can choose to compile this image yourself as follows:

touch Dockerfile
# Copy the above Dockerfile content into the Dockerfile, then
docker build -t my-nvidia .

Or you can pull this base image with the following command:

docker pull hub.aiursoft.cn/aiursoft/internalimages/nvidia:latest

3. Configure the SwinIR Container

Next, we will configure a SwinIR container based on the base image and install the dependencies for SwinIR.

FROM my-nvidia
RUN sudo apt update && sudo apt install -y git
RUN git clone https://github.com/JingyunLiang/SwinIR.git

RUN pip install opencv-python torch torchvision
RUN pip install requests numpy timm

# Download the model
RUN mkdir -p model_zoo/swinir && \
    wget -O model_zoo/swinir/003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN.pth https://github.com/JingyunLiang/SwinIR/releases/download/v0.0/003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN.pth

RUN mkdir -p /app/input
RUN mkdir -p /app/results/swinir_real_sr_x4_large

VOLUME /app/input
VOLUME /app/results/swinir_real_sr_x4_large

ENTRYPOINT ["python3", "./SwinIR/main_test_swinir.py", "--task", "real_sr", "--scale", "4", "--large_model", "--model_path", "model_zoo/swinir/003_realSR_BSRGAN_DFOWMFC_s64w8_SwinIR-L_x4_GAN.pth", "--folder_lq", "/app/input", "--tile", "192", "--tile_overlap", "16"]

# Usage
# Prepare two folders: input and output
# sudo docker run --gpus all -it -v $(pwd)/input:/app/input -v $(pwd)/output:/app/results/swinir_real_sr_x4_large tagname

This container will download the SwinIR code and model and configure it to run directly. You can build this container with the following command:

docker build -t swinir .


4. Run the Container

Next, let's try using the image you just compiled to perform super-resolution reconstruction on an image.

In the ~/Demo folder, I created input and output folders. I placed an image in the input folder.


Next, in the ~/Demo folder, run the following command:

sudo docker run --gpus all -it -v $(pwd)/input:/app/input -v $(pwd)/output:/app/results/swinir_real_sr_x4_large swinir


This command will run your container, mount the input folder to the container's /app/input folder, and the output folder to the container's /app/results/swinir_real_sr_x4_large folder.

After successfully running, the images in the input folder will be super-resolved and saved in the output folder.




5. Conclusion

This is the complete process for using SwinIR for image super-resolution reconstruction. With this technology, you can restore old photos, enhance low-resolution images, and even improve surveillance videos.