DDColor: Photo-Realistic Image Colorization

To launch a web-based interactive demo, install gradio and gradio_imageslider, then run the Gradio application script.

pip install gradio gradio_imageslider
python demo/gradio_app.py

To set up DDColor in a clean environment, it is recommended to use Conda. This process involves creating a Python 3.9 environment, installing PyTorch with CUDA support, and then installing the project requirements.

Requirements:

• Python >= 3.7
• PyTorch >= 1.7

For Training: If you intend to train the model, you must also install additional dependencies and basicsr using the training requirements file and running the setup script in develop mode.

conda create -n ddcolor python=3.9
conda activate ddcolor
pip install torch==2.2.0 torchvision==0.17.0 --index-url https://download.pytorch.org/whl/cu118

pip install -r requirements.txt

# For training, install the following additional dependencies and basicsr
pip install -r requirements.train.txt
python3 setup.py develop

You can perform image colorization locally without requiring basicsr.

1. Download the pretrained model using ModelScope's snapshot download.
2. Run the inference script by pointing to the downloaded model path and your input image directory.

from modelscope.hub.snapshot_download import snapshot_download

model_dir = snapshot_download('damo/cv_ddcolor_image-colorization', cache_dir='./modelscope')
print('model assets saved to %s' % model_dir)

python scripts/infer.py --model_path ./modelscope/damo/cv_ddcolor_image-colorization/pytorch_model.pt --input ./assets/test_images

Training DDColor involves several steps:

1. Dataset Preparation: Use python scripts/get_meta_file.py to generate a dataset list file from ImageNet or a custom dataset.
2. Pretrained Weights: Download weights for ConvNeXt and InceptionV3 and place them in the pretrain folder.
3. Configuration: Edit options/train/train_ddcolor.yml to specify the meta_info_file and other training hyperparameters.
4. Execution: Run the training script via sh scripts/train.sh.

DDColor supports ONNX export. You must first install onnx, onnxruntime, and onnxsim.

Usage: Run the export_onnx.py script providing the path to your pretrained .pth model and the desired output path for the .onnx file.

pip install onnx==1.16.1 onnxruntime==1.19.2 onnxsim==0.4.36

python scripts/export_onnx.py --model_path pretrain/ddcolor_paper_tiny.pth --export_path weights/ddcolor-tiny.onnx

You can use the ModelScope pipeline API to perform colorization. This method automatically handles model downloading and inference. Ensure you have modelscope and opencv-python installed.

import cv2
from modelscope.outputs import OutputKeys
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

img_colorization = pipeline(Tasks.image_colorization, model='damo/cv_ddcolor_image-colorization')
result = img_colorization('https://modelscope.oss-cn-beijing.aliyuncs.com/test/images/audrey_hepburn.jpg')
cv2.imwrite('result.png', result[OutputKeys.OUTPUT_IMG])

DDColor provides several pre-trained models optimized for different use cases. When performing inference, choose a model_name based on your requirements:

• ddcolor_modelscope (Default): Best for general use and testing images outside of ImageNet. It uses a specific data cleaning scheme to achieve high qualitative results with minimal FID degradation.
• ddcolor_artistic: Best for high-quality artistic results. It was trained on an extended dataset of artistic images and does not use colorfulness loss, which may reduce color artifacts.
• ddcolor_paper: Use this only if you need to reproduce specific images from the original DDColor paper (DDColor-L trained on ImageNet).
• ddcolor_paper_tiny: The most lightweight version (DDColor-T trained on ImageNet).

To measure the performance of your DDColor ONNX deployment, run the inference session multiple times (e.g., 100 iterations) and calculate the average time per batch, per image, and the resulting FPS (Frames Per Second).

import time
import tqdm
import numpy as np

# Run the model 100 times to get an average time
times = []
for i in tqdm.tqdm(range(100)):
    start = time.time()
    outputs = ort_session.run(None, inputs)
    times.append(time.time() - start)

print(f"Average time per batch: {np.mean(times):.4f} seconds")
print(f"Average time per image: {np.mean(times)/batch_size:.4f} seconds")
print(f"Average FPS per image: {batch_size/np.mean(times):.4f}")